Update test.py

test.py

@@ -1,9 +1,10 @@
+import time
+import torch
 from transformers import T5ForConditionalGeneration, AutoTokenizer
 from datasets import load_dataset
-import torch
 
 # Путь к модели и данным
-model_path = "./flan-t5-autobatch"
+model_path = "./unzipped_model"
 validation_file = "mt5_validation_data-1.jsonl"
 
 # Загрузка модели и токенизатора
@@ -11,9 +12,9 @@ tokenizer = AutoTokenizer.from_pretrained(model_path)
 model = T5ForConditionalGeneration.from_pretrained(model_path)
 model.eval()
 
-# Используем GPU если есть
+# Используем GPU, если доступно
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-model = model.to(device)
+model.to(device)
 
 # Загрузка валидационной выборки
 dataset = load_dataset("json", data_files={"validation": validation_file})
@@ -21,7 +22,14 @@ val_data = dataset["validation"]
 
 # Функция предсказания
 def predict(text):
-    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True, max_length=256).to(device)
+    inputs = tokenizer(
+        text,
+        return_tensors="pt",
+        truncation=True,
+        padding=True,
+        max_length=256
+    ).to(device)
+    
     outputs = model.generate(
         **inputs,
         max_length=64,
@@ -30,55 +38,74 @@ def predict(text):
     )
     return tokenizer.decode(outputs[0], skip_special_tokens=True)
 
-# Подсчёт точности
-correct = 0
-correct_country = 0
-correct_city = 0
-results = []
+# Счётчики
+country_correct = 0
+city_correct = 0
+full_correct = 0
 
-for idx, example in enumerate(val_data):
-    text = example["text"]
-    target = example["target"].strip()
-    pred = predict(text).strip()
+# Массив для хранения ошибочных предсказаний
+# (где хотя бы что-то не совпало между таргетом и предиктом)
+incorrect_samples = []
 
-    # Полное совпадение
-    if pred == target:
-        correct += 1
+start_time = time.time()
+
+for example in val_data:
+    text = example["text"]
+    target = example["target"]
+    
+    # Получаем предсказание
+    prediction = predict(text)
+    
+    # Разбиваем на "Страна:Город" - если нет двоеточия, ставим "unknown" для второй части
+    pred_parts = prediction.split(":", 1)
+    if len(pred_parts) == 2:
+        pred_country, pred_city = pred_parts[0].strip(), pred_parts[1].strip()
+    else:
+        pred_country = pred_parts[0].strip()
+        pred_city = "unknown"
 
-    # Приведение к нижнему регистру и очистка пробелов
-    target_parts = [x.strip().lower() for x in target.split(",")]
-    pred_parts = [x.strip().lower() for x in pred.split(",")]
+    # Аналогично таргет
+    target_parts = target.split(":", 1)
+    if len(target_parts) == 2:
+        true_country, true_city = target_parts[0].strip(), target_parts[1].strip()
+    else:
+        true_country = target_parts[0].strip()
+        true_city = "unknown"
 
-    # Сравнение страны (первый элемент)
-    if target_parts and pred_parts and target_parts[0] == pred_parts[0]:
-        correct_country += 1
+    # Сравниваем
+    if pred_country == true_country:
+        country_correct += 1
+    if pred_city == true_city:
+        city_correct += 1
+    if (pred_country == true_country) and (pred_city == true_city):
+        full_correct += 1
+    else:
+        # Если хотя бы что-то не совпало, сохраним этот пример
+        incorrect_samples.append({
+            "text": text,
+            "target": f"{true_country}:{true_city}",
+            "prediction": f"{pred_country}:{pred_city}"
+        })
 
-    # Сравнение города (только если в 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
+end_time = time.time()
+total_time = end_time - start_time
 
-    results.append({
-        "text": text,
-        "pred": pred,
-        "target": target,
-        "match": pred == target
-    })
+num_examples = len(val_data)
+time_per_example = total_time / num_examples if num_examples > 0 else 0
 
-# Вывод до 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"📝 Вход: {r['text']}")
-    print(f"✅ Target: {r['target']}")
-    print(f"🤖 Предсказание: {r['pred']}")
-    print("-" * 50)
+# Выводим первые 80 «ошибочных» примеров (если их меньше, то все)
+print("Приме��ы, где хотя бы что-то не совпало (макс. 80):")
+for i, item in enumerate(incorrect_samples[:80]):
+    print(f"\nПример {i+1}:")
+    print(f"Текст:        {item['text']}")
+    print(f"Таргет:       {item['target']}")
+    print(f"Предсказание: {item['prediction']}")
 
-# Accuracy
-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")
+# После этого выводим статистику
+print("\nРезультаты:")
+print(f"Всего примеров валидации: {num_examples}")
+print(f"Совпало стран: {country_correct}")
+print(f"Совпало городов: {city_correct}")
+print(f"Полных совпадений (страна и город): {full_correct}")
+print(f"Общее время выполнения скрипта: {total_time:.4f} сек.")
+print(f"Время на одно предсказание: {time_per_example:.6f} сек.")