Bakalarska_praca/metric_run.py

143 lines
8.5 KiB
Python
Raw Normal View History

2024-10-29 13:08:43 +00:00
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}")