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

added temperature and token count to ai config

Browse Source
This commit is contained in:
wea_ondara
2024-05-29 18:36:17 +02:00
parent 3a78660883
commit bc0712ba11
14 changed files with 197 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
ai/QwenAi.py
View File

@@ -13,7 +13,7 @@ class QwenAi(AiBase):
def __init__(self, model_id_or_path=default_model_id):
super().__init__(model_id_or_path)
def generate(self, messages):
def generate(self, messages, token_count=100, **kwargs):
try:
# prepare
messages = [m for m in messages if m['role'] != 'system']
@@ -22,7 +22,7 @@ class QwenAi(AiBase):
# generate
text = self.tokenizer.apply_chat_template(input_messages, tokenize=False, add_generation_prompt=True)
model_inputs = self.tokenizer([text], return_tensors='pt').to(self.default_device)
generated_ids = self.model.generate(model_inputs.input_ids, max_new_tokens=300)
generated_ids = self.model.generate(model_inputs.input_ids, max_new_tokens=token_count)
generated_ids = [
output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
]