diff --git a/sia/__main__.py b/sia/__main__.py index 04eccd9..7ff3842 100644 --- a/sia/__main__.py +++ b/sia/__main__.py @@ -67,7 +67,6 @@ class Main: self._llms['qwq'] = QwQLlmEngine( config.qwq_model, config.qwq_temperature, - config.qwq_token_limit, self._action_schema, ) diff --git a/sia/llm_engine/qwq_llm_engine.py b/sia/llm_engine/qwq_llm_engine.py index 9f29fe7..d8770ad 100644 --- a/sia/llm_engine/qwq_llm_engine.py +++ b/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 \ No newline at end of file + return self._pipeline.model.config.max_position_embeddings \ No newline at end of file