Update app.py

app.py

@@ -6,8 +6,7 @@ examples = [
     "Caesar in Gallia bellum gessit. <mask> instructae erant, sed subito",
     "Est autem et aliud genus testudinis, quod <mask>, quemadmodum quae supra scripta sunt",
     "Quemadmodum vero minores rotae duriores et <mask>, sic phalangae et iuga, in quibus partibus habent minora a centro ad capita intervalla",
-    "illud additur, ne, qui certum ordinem ex <mask>, ulli vos alteri hominum generi haerere vereamini nec timeatis vos",
-    "Gli italiani perdono le partite di calcio come se <mask> e perdono le guerre come se fossero partite di calcio"
+    "illud additur, ne, qui certum ordinem ex <mask>, ulli vos alteri hominum generi haerere vereamini nec timeatis vos"
 ]
 examples_correct = [
     "Omnes legiones",
@@ -53,9 +52,13 @@ input_text = st.text_area(
 model_name = "Cicciokr/BART-la-s"
 model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
 tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=False)
+tokenizer.pad_token_id = tokenizer.eos_token_id
+
+model_name_mio = "Cicciokr/BART-CC100-la"
+model_mio = AutoModelForSeq2SeqLM.from_pretrained(model_name_mio)
+tokenizer_mio = AutoTokenizer.from_pretrained(model_name_mio, use_fast=False)
+tokenizer_mio.pad_token_id = tokenizer_mio.eos_token_id
 
-#model_name_mio = "Cicciokr/mbart50-large-latin"
-#tokenizer.pad_token_id = tokenizer.eos_token_id
 #generator = pipeline("text2text-generation", model=model, tokenizer=tokenizer)
 #generator_mio = pipeline("text2text-generation", model=model_mio, tokenizer=tokenizer_mio)
 
@@ -67,16 +70,30 @@ if input_text:
     output = model.generate(
         **inputs,
         max_length=512,
-    #    num_beams=4,
-    #    num_return_sequences=1,
+        num_beams=4,
+        num_return_sequences=1,
         do_sample=True,
         temperature=0.9
-    #    top_k=1,
+        top_k=1,
     )
     generated_text = tokenizer.decode(output[0], skip_special_tokens=True)
     if 'input_text_value_correct' in st.session_state:
         st.write(f" Parola corretta: {st.session_state['input_text_value_correct']}\n")
-    st.subheader("Risultato BART:")
+    st.subheader("Risultato BART TheLatinLibrary:")
     st.write(f" Frase predetta: {generated_text}\n")
     #st.write(f" Frase predetta: {tokenizer.decode(output[0], skip_special_tokens=True)}\n")
     #print(output)
+
+    inputs_mio = tokenizer_mio(input_text, return_tensors="pt")
+    output_mio = model_mio.generate(
+        **inputs_mio,
+        max_length=512,
+        num_beams=4,
+        num_return_sequences=1,
+        do_sample=True,
+        temperature=0.9
+        top_k=1,
+    )
+    generated_text_mio = tokenizer_mio.decode(output[0], skip_special_tokens=True)
+    st.subheader("Risultato BART CC100:")
+    st.write(f" Frase predetta: {generated_text_mio}\n")