Update test.py

test.py

@@ -1,14 +1,14 @@
-from transformers import MT5ForConditionalGeneration, MT5Tokenizer
+from transformers import T5ForConditionalGeneration, AutoTokenizer
 from datasets import load_dataset
 import torch
 
 # Путь к модели и данным
-model_path = "./mt5-finetuned"
+model_path = "./flan-t5-autobatch"
 validation_file = "mt5_validation_data-1.jsonl"
 
 # Загрузка модели и токенизатора
-tokenizer = MT5Tokenizer.from_pretrained(model_path)
-model = MT5ForConditionalGeneration.from_pretrained(model_path)
+tokenizer = AutoTokenizer.from_pretrained(model_path)
+model = T5ForConditionalGeneration.from_pretrained(model_path)
 model.eval()
 
 # Используем GPU если есть
@@ -32,6 +32,8 @@ def predict(text):
 
 # Подсчёт точности
 correct = 0
+correct_country = 0
+correct_city = 0
 results = []
 
 for idx, example in enumerate(val_data):
@@ -39,19 +41,44 @@ for idx, example in enumerate(val_data):
     target = example["target"].strip()
     pred = predict(text).strip()
 
-    results.append((text, pred, target))
+    # Полное совпадение
     if pred == target:
         correct += 1
 
-# Примеры
-print("📋 Примеры предсказаний:\n")
-for i, (text, pred, target) in enumerate(results[:80]): # кол-во примеров 
+    # Приведение к нижнему регистру и очистка пробелов
+    target_parts = [x.strip().lower() for x in target.split(",")]
+    pred_parts = [x.strip().lower() for x in pred.split(",")]
+
+    # Сравнение страны (первый элемент)
+    if target_parts and pred_parts and target_parts[0] == pred_parts[0]:
+        correct_country += 1
+
+    # Сравнение города (только если в target есть второй элемент)
+    if len(target_parts) > 1:
+        target_city = target_parts[-1]
+        pred_city = pred_parts[-1] if len(pred_parts) > 1 else ""
+        if target_city == pred_city:
+            correct_city += 1
+
+    results.append({
+        "text": text,
+        "pred": pred,
+        "target": target,
+        "match": pred == target
+    })
+
+# Вывод до 80 первых НЕСОВПАВШИХ
+mismatched = [r for r in results if not r["match"]]
+print("📋 Несовпавшие предсказания (до 80):\n")
+for i, r in enumerate(mismatched[:80]):
     print(f"#{i+1}")
-    print(f"📝 Вход: {text}")
-    print(f"✅ Target: {target}")
-    print(f"🤖 Предсказание: {pred}")
+    print(f"📝 Вход: {r['text']}")
+    print(f"✅ Target: {r['target']}")
+    print(f"🤖 Предсказание: {r['pred']}")
     print("-" * 50)
 
 # Accuracy
-accuracy = correct / len(val_data)
-print(f"\n✅ Accuracy: {accuracy:.4f} ({correct}/{len(val_data)})\n")
+total = len(val_data)
+print(f"\n✅ Accuracy (полное совпадение): {correct / total:.4f} ({correct}/{total})")
+print(f"🌍 Accuracy по странам: {correct_country / total:.4f} ({correct_country}/{total})")
+print(f"🏙️ Accuracy по городам: {correct_city / total:.4f} ({correct_city}/{total})\n")