143 lines
8.5 KiB
Python
143 lines
8.5 KiB
Python
import torch
|
||
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
|
||
import jiwer
|
||
import sacrebleu
|
||
|
||
# 1. Загрузка модели и токенизатора
|
||
model_path = "google/mt5-large" # Укажите путь к вашей модели
|
||
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}")
|
||
|