Update fine_tune_inference_test_mistral.py

fine_tune_inference_test_mistral.py

@@ -76,10 +76,11 @@ def chat(msg: Message):
             return {"error": "Boş giriş"}
 
         messages = [{"role": "user", "content": user_input}]
-        input_ids = tokenizer.apply_chat_template(messages, return_tensors="pt", add_generation_prompt=True).to(model.device)
+        inputs = tokenizer.apply_chat_template(messages, return_tensors="pt", add_generation_prompt=True)
+        inputs = {k: v.to(model.device) for k, v in inputs.items()}
 
         generate_args = {
-            "max_new_tokens": 1024,
+            "max_new_tokens": 512,
             "return_dict_in_generate": True,
             "output_scores": True,
             "do_sample": USE_SAMPLING
@@ -93,10 +94,11 @@ def chat(msg: Message):
             })
 
         with torch.no_grad():
-            output = model.generate(input_ids=input_ids, **generate_args)
+            output = model.generate(**inputs, **generate_args)
 
+        prompt_text = tokenizer.decode(inputs["input_ids"][0], skip_special_tokens=True)
         decoded = tokenizer.decode(output.sequences[0], skip_special_tokens=True)
-        answer = decoded.split("</s>")[-1].strip()
+        answer = decoded.replace(prompt_text, "").strip()
 
         if output.scores and len(output.scores) > 0:
             first_token_score = output.scores[0][0]