Minor tweaks to the tensors loading

eval_utils.py

@@ -11,7 +11,7 @@ def predict_top_p(model, tokenizer, text, P=0.9):
     """
     assert P <= 1 and P >= 0
 
-    logits = model(**tokenizer(text, return_tensors="pt")).logits
+    logits = model(**tokenizer(text, return_tensors="pt").to(device)).logits
     probabilities = torch.softmax(logits, dim=1).flatten().tolist()
     topk_predictions = torch.topk(logits, 18).indices.flatten().tolist()
 
@@ -71,12 +71,11 @@ def prompt_chat_LLM(model, tokenizer, text):
             },
         ]
         input_ids = tokenizer.apply_chat_template(
-            messages, tokenize=True, add_generation_prompt=True, return_tensors="pt", device_map=device
-        )
-        gen_tokens = model.generate(input_ids, max_new_tokens=20)
+            messages, tokenize=True, add_generation_prompt=True, return_dict=True, return_tensors="pt").to("cuda")
+        gen_tokens = model.generate(**input_ids, max_new_tokens=20, pad_token_id=tokenizer.eos_token_id)
         gen_text = tokenizer.decode(gen_tokens[0], skip_special_tokens=True)
         # TODO: Add a condition for the case of "لا" and other responses (e.g., refuse to answer)
-        if "نعم" in gen_text:
+        if gen_text.strip(". ").endswith("نعم"):
             predicted_dialects.append(dialect)
     return predicted_dialects