Fixed qwq token limit

sia/llm_engine/qwq_llm_engine.py

@@ -1,7 +1,7 @@
 from pathlib import Path
 from threading import Thread
 from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline, TextIteratorStreamer, BitsAndBytesConfig
-from typing import Callable, Iterator, Optional, List
+from typing import Callable, Iterator, Optional
 from xml_schema_validator import XmlLogitsProcessor
 import json
 import torch
@@ -15,7 +15,6 @@ class QwQLlmEngine(LlmEngine):
         self,
         model_path: Path,
         temperature: float,
-        token_limit: int = None,
         xml_schema_text: Optional[str] = None,
     ):
         """
@@ -24,11 +23,9 @@ class QwQLlmEngine(LlmEngine):
         Args:
             model_path: Local path to the model
             temperature: Sampling temperature
-            token_limit: Maximum tokens to generate
             xml_schema_text: Optional XML schema to validate against
         """
         self._temperature = temperature
-        self._token_limit = token_limit
 
         with open('/root/sia/qwq_tokenizer_config.json', 'r') as f:
             tokenizer_config = json.load(f)
@@ -53,7 +50,7 @@ class QwQLlmEngine(LlmEngine):
             tokenizer_config=tokenizer_config,
         )
 
-        self._pipline = pipeline(
+        self._pipeline = pipeline(
             "text-generation",
             model=model,
             tokenizer=self._tokenizer,
@@ -98,7 +95,7 @@ class QwQLlmEngine(LlmEngine):
             "text_inputs": text,
             "do_sample": True,
             "temperature": self._temperature,
-            "max_new_tokens": self._token_limit,
+            "max_new_tokens": self.token_limit(),
             "streamer": streamer,
         }
         
@@ -106,7 +103,7 @@ class QwQLlmEngine(LlmEngine):
             generation_kwargs["logits_processor"] = [self._logits_processor.copy()]
 
         generation_thread = Thread(
-            target=self._pipline,
+            target=self._pipeline,
             kwargs=generation_kwargs
         )
 
@@ -140,4 +137,4 @@ class QwQLlmEngine(LlmEngine):
         return len(self._tokenizer.encode(prompt))
 
     def token_limit(self) -> int:
-        return self._token_limit
+        return self._pipeline.model.config.max_position_embeddings