baseline app

app.py

@@ -5,28 +5,15 @@ from datasets import load_dataset
 # Załaduj dataset z Hugging Face
 dataset = load_dataset("clui/lyricsgenius")
 
-# Funkcja do ładowania modeli tłumaczeniowych z cache
-@st.cache_resource
-def load_translation_models():
-    model_name_it_en = "Helsinki-NLP/opus-mt-it-en"
-    tokenizer_it_en = AutoTokenizer.from_pretrained(model_name_it_en)
-    model_it_en = AutoModelForSeq2SeqLM.from_pretrained(model_name_it_en)
+# Wybór modelu do tłumaczenia z włoskiego na angielski
+model_name_it_en = "Helsinki-NLP/opus-mt-it-en"
+tokenizer_it_en = AutoTokenizer.from_pretrained(model_name_it_en)
+model_it_en = AutoModelForSeq2SeqLM.from_pretrained(model_name_it_en)
 
-    model_name_en_pl = "gsarti/opus-mt-tc-en-pl"
-    tokenizer_en_pl = AutoTokenizer.from_pretrained(model_name_en_pl)
-    model_en_pl = AutoModelForSeq2SeqLM.from_pretrained(model_name_en_pl)
-
-    return tokenizer_it_en, model_it_en, tokenizer_en_pl, model_en_pl
-
-# Załaduj modele tłumaczeniowe
-tokenizer_it_en, model_it_en, tokenizer_en_pl, model_en_pl = load_translation_models()
-
-# Funkcja do tłumaczenia tekstu
-def translate(text, tokenizer, model):
-    input_ids = tokenizer(text, return_tensors="pt", padding=True).input_ids
-    outputs = model.generate(input_ids)
-    translated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return translated_text
+# Wybór modelu do tłumaczenia z angielskiego na polski
+model_name_en_pl = "sdadas/mt5-base-translator-en-pl"
+tokenizer_en_pl = AutoTokenizer.from_pretrained(model_name_en_pl)
+model_en_pl = AutoModelForSeq2SeqLM.from_pretrained(model_name_en_pl)
 
 # Pole do wprowadzania nazwy artysty
 st.title("Neural Notes")
@@ -68,7 +55,9 @@ if artist_name:
 
                     for line in lines:
                         if line.strip():
-                            translated_text = translate(line, tokenizer_it_en, model_it_en)
+                            inputs = tokenizer_it_en(line, return_tensors="pt", padding=True)
+                            translated_outputs = model_it_en.generate(**inputs)
+                            translated_text = tokenizer_it_en.decode(translated_outputs[0], skip_special_tokens=True)
                             translated_lines_en.append(translated_text)
                         else:
                             translated_lines_en.append("")
@@ -84,7 +73,9 @@ if artist_name:
 
                     for line in translated_lines_en:
                         if line.strip():
-                            translated_text = translate(line, tokenizer_en_pl, model_en_pl)
+                            inputs = tokenizer_en_pl(line, return_tensors="pt", padding=True)
+                            translated_outputs = model_en_pl.generate(**inputs)
+                            translated_text = tokenizer_en_pl.decode(translated_outputs[0], skip_special_tokens=True)
                             translated_lines_pl.append(translated_text)
                         else:
                             translated_lines_pl.append("")