API fix

models.py

@@ -1,13 +1,19 @@
 import torch
-from transformers import pipeline
+from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer
 
 def get_tinyllama():
     tinyllama = pipeline("text-generation", model="TinyLlama/TinyLlama-1.1B-Chat-v1.0", torch_dtype=torch.float16, device_map="auto")
     return tinyllama
 
 def get_qwen2ins1b():
-    tinyllama = pipeline("text-generation", model="Qwen/Qwen2-1.5B-Instruct", torch_dtype=torch.float16, device_map="auto")
-    return tinyllama
+    model = AutoModelForCausalLM.from_pretrained(
+        "Qwen/Qwen2-1.5B-Instruct",
+        torch_dtype="auto",
+        device_map="auto"
+        )
+    
+    tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2-1.5B-Instruct")
+    return {'model': model, 'tokenizer': tokenizer}
 
 def response_tinyllama(
         model=None,
@@ -46,7 +52,21 @@ def response_qwen2ins1b(
         if len(step) >= 2:
             messages_dict.append({'role': 'assistant', 'content': step[1]})
 
-    prompt = model.tokenizer.apply_chat_template(messages_dict, tokenize=False, add_generation_prompt=True)
-    outputs = model(prompt, max_new_tokens=64, do_sample=True, temperature=0.7, top_k=50, top_p=0.95)
+    text = model['tokenizer'].apply_chat_template(
+        messages_dict,
+        tokenize=False,
+        add_generation_prompt=True
+    )
+    model_inputs = model['tokenizer']([text], return_tensors="pt")
+
+    generated_ids = model['model'].generate(
+        model_inputs.input_ids,
+        max_new_tokens=512
+    )
+    generated_ids = [
+        output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
+    ]
+
+    response = model['tokenizer'].batch_decode(generated_ids, skip_special_tokens=True)[0]
 
-    return outputs[0]['generated_text'] #.split('<|assistant|>')[1].strip()
+    return response # outputs[0]['generated_text'] #.split('<|assistant|>')[1].strip()