import torch
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
import jiwer
import sacrebleu

# 1. Загрузка модели и токенизатора
model_path = "T5Autocorrection1337"  # Укажите путь к вашей модели
tokenizer = AutoTokenizer.from_pretrained(model_path)
model = AutoModelForSeq2SeqLM.from_pretrained(model_path)

# Устройство (GPU/CPU)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)

# 2. Функция для генерации предсказаний модели
def generate_correction(sentence):
    inputs = tokenizer(sentence, return_tensors="pt", max_length=128, truncation=True, padding="max_length").to(device)
    with torch.no_grad():
        outputs = model.generate(inputs["input_ids"], max_length=128, num_beams=5, early_stopping=True)
    corrected_sentence = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return corrected_sentence

# 3. Пример массива с тестовыми данными для оценки метрик
test_data = [
    {"incorrect": "On chodi do skoly kazdy den", "correct": "On chodí do školy každý deň"},
    {"incorrect": "Rano vstavam o siestej hodine", "correct": "Ráno vstávam o šiestej hodine"},
    {"incorrect": "Toto je pekny dom", "correct": "Toto je pekný dom"},
    {"incorrect": "Moja sestra studuje na univerzite", "correct": "Moja sestra študuje na univerzite"},
    {"incorrect": "Ideme do obchodu kupit chleba", "correct": "Ideme do obchodu kúpiť chlieb"},
    {"incorrect": "Mam rad knizky o histórii", "correct": "Mám rád knižky o histórii"},
    {"incorrect": "Včera som bol v kine", "correct": "Včera som bol v kine"},
    {"incorrect": "To je moj oblubeny film", "correct": "To je môj obľúbený film"},
    {"incorrect": "V nedelu ideme na vylety", "correct": "V nedeľu ideme na výlety"},
    {"incorrect": "Moji rodicia su velmi dobri ludia", "correct": "Moji rodičia sú veľmi dobrí ľudia"},
    
    {"incorrect": "Tvoj brat je starsi ako tvoja sestra", "correct": "Tvoj brat je starší ako tvoja sestra"},
    {"incorrect": "Slnko svieti kazdy den", "correct": "Slnko svieti každý deň"},
    {"incorrect": "On studuje chemiju na univerzite", "correct": "On študuje chémiu na univerzite"},
    {"incorrect": "Mama vari veceru", "correct": "Mama varí večeru"},
    {"incorrect": "V pondelok mame volny den", "correct": "V pondelok máme voľný deň"},
    {"incorrect": "Oni hovorili o historii mesta", "correct": "Oni hovorili o histórii mesta"},
    {"incorrect": "On bol moj oblubeny ucitel", "correct": "On bol môj obľúbený učiteľ"},
    {"incorrect": "V piatok sa stretneme s priatelmi", "correct": "V piatok sa stretneme s priateľmi"},
    {"incorrect": "Moja sestra spieva krasne piesne", "correct": "Moja sestra spieva krásne piesne"},
    {"incorrect": "Oni cestuju po celom svete", "correct": "Oni cestujú po celom svete"},

    {"incorrect": "V lese je vela stromov", "correct": "V lese je veľa stromov"},
    {"incorrect": "Stary clovek rozprava pomaly", "correct": "Starý človek rozpráva pomaly"},
    {"incorrect": "Ideme sa prechadzat po parku", "correct": "Ideme sa prechádzať po parku"},
    {"incorrect": "V nedelu budu prsat", "correct": "V nedeľu bude pršať"},
    {"incorrect": "Maly chlapec sa hra s loptou", "correct": "Malý chlapec sa hrá s loptou"},
    {"incorrect": "Otec pracuje v zahrade", "correct": "Otec pracuje v záhrade"},
    {"incorrect": "Moja babka rada pecie kolace", "correct": "Moja babka rada pečie koláče"},
    {"incorrect": "Sused ma nove auto", "correct": "Sused má nové auto"},
    {"incorrect": "On ma vela prace dnes", "correct": "On má veľa práce dnes"},
    {"incorrect": "V skole sa uci matematika a fyzika", "correct": "V škole sa učí matematika a fyzika"},

    {"incorrect": "Ja mam rada zeleninu", "correct": "Ja mám rada zeleninu"},
    {"incorrect": "On pise list svojej sestre", "correct": "On píše list svojej sestre"},
    {"incorrect": "Stary dom stoji v strede mesta", "correct": "Starý dom stojí v strede mesta"},
    {"incorrect": "V sobotu mame party", "correct": "V sobotu máme párty"},
    {"incorrect": "Oni kupuju novy nabytok", "correct": "Oni kupujú nový nábytok"},
    {"incorrect": "On studuje informatiku", "correct": "On študuje informatiku"},
    {"incorrect": "Moja mama pripravuje obed", "correct": "Moja mama pripravuje obed"},
    {"incorrect": "On ma dve knihy", "correct": "On má dve knihy"},
    {"incorrect": "Vecer budeme pozerat film", "correct": "Večer budeme pozerať film"},
    {"incorrect": "Moje dieta je velmi sikovne", "correct": "Moje dieťa je veľmi šikovné"},

    {"incorrect": "Oni budu cestovat do Pariza", "correct": "Oni budú cestovať do Paríža"},
    {"incorrect": "Mam doma dve macky", "correct": "Mám doma dve mačky"},
    {"incorrect": "On plava v bazene kazdy den", "correct": "On pláva v bazéne každý deň"},
    {"incorrect": "V tejto knihe je vela obrazkov", "correct": "V tejto knihe je veľa obrázkov"},
    {"incorrect": "Stary clovek si cita noviny", "correct": "Starý človek si číta noviny"},
    {"incorrect": "My budeme oslavovat narodeniny", "correct": "My budeme oslavovať narodeniny"},
    {"incorrect": "Moja sestra ma rada jablka", "correct": "Moja sestra má rada jablká"},
    {"incorrect": "Ideme na vylet do hor", "correct": "Ideme na výlet do hôr"},
    {"incorrect": "On mal rad svoju pracu", "correct": "On mal rád svoju prácu"},
    {"incorrect": "Oni si kupili nove knihy", "correct": "Oni si kúpili nové knihy"},

    {"incorrect": "Pocasie je dnes velmi dobre", "correct": "Počasie je dnes veľmi dobré"},
    {"incorrect": "Maly chlapec je velmi radostny", "correct": "Malý chlapec je veľmi radostný"},
    {"incorrect": "Mama pripravuje veceru", "correct": "Mama pripravuje večeru"},
    {"incorrect": "V parku je vela deti", "correct": "V parku je veľa detí"},
    {"incorrect": "Ja mam dve sestry", "correct": "Ja mám dve sestry"},
    {"incorrect": "On je velmi inteligentny student", "correct": "On je veľmi inteligentný študent"},
    {"incorrect": "Moja babka ma rada rozpravky", "correct": "Moja babka má rada rozprávky"},
    {"incorrect": "Oni sa stretnu o piatej hodine", "correct": "Oni sa stretnú o piatej hodine"},
    {"incorrect": "On cita knihu o historii", "correct": "On číta knihu o histórii"},
    {"incorrect": "Ja mam rad citanie knih", "correct": "Ja mám rád čítanie kníh"},

    {"incorrect": "Oni maju nove auto", "correct": "Oni majú nové auto"},
    {"incorrect": "Moja mama vari polievku", "correct": "Moja mama varí polievku"},
    {"incorrect": "On ma rad matematiku", "correct": "On má rád matematiku"},
    {"incorrect": "Maly chlapec sa hra na dvore", "correct": "Malý chlapec sa hrá na dvore"},
    {"incorrect": "Vecer budeme pozerat televiziu", "correct": "Večer budeme pozerať televíziu"},
    {"incorrect": "V meste je nova kaviaren", "correct": "V meste je nová kaviareň"},
    {"incorrect": "Moj brat studuje na univerzite", "correct": "Môj brat študuje na univerzite"},
    {"incorrect": "On ma rad svoju pracu", "correct": "On má rád svoju prácu"},
    {"incorrect": "V lese je vela stromov", "correct": "V lese je veľa stromov"},
    {"incorrect": "Mama pripravuje obed", "correct": "Mama pripravuje obed"}
]


# Списки для хранения истинных и предсказанных значений
references = []
predictions = []

# 4. Генерация предсказаний для тестового набора
for item in test_data:
    incorrect_sentence = item['incorrect']
    correct_sentence_reference = item['correct']
    
    # Генерация предсказания
    prediction = generate_correction(incorrect_sentence)  # Вызов функции generate_correction
    references.append(correct_sentence_reference)
    predictions.append(prediction)

# 5. Оценка метрик

# WER (Word Error Rate)
wer = jiwer.wer(references, predictions)
print(f"WER: {wer}")

# CER (Character Error Rate)
cer = jiwer.cer(references, predictions)
print(f"CER: {cer}")

# SER (Sentence Error Rate)
def calculate_ser(references, predictions):
    errors = 0
    for ref, pred in zip(references, predictions):
        if ref != pred:
            errors += 1
    ser = errors / len(references)
    return ser

ser = calculate_ser(references, predictions)
print(f"SER: {ser}")
# BLEU (Bilingual Evaluation Understudy Score)
bleu = sacrebleu.corpus_bleu(predictions, [references])
print(f"BLEU: {bleu.score}")