llm/SIA
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Fixed qwq token limit

Browse Source
This commit is contained in:
Geens
2025-04-11 18:03:17 +02:00
parent e9f5bcc519
commit 7e254b7f7e
2 changed files with 5 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
sia/__main__.py
View File

@@ -67,7 +67,6 @@ class Main:
self._llms['qwq'] = QwQLlmEngine( self._llms['qwq'] = QwQLlmEngine(
config.qwq_model, config.qwq_model,
config.qwq_temperature, config.qwq_temperature,
config.qwq_token_limit,
self._action_schema, self._action_schema,
) )

13
sia/llm_engine/qwq_llm_engine.py
View File

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