Switch PAWS by XNLI

app.py

@@ -34,20 +34,11 @@ PROMPT_LIST = [
 
 PAWS_X_PROMPT_LIST = [
     "Te amo.</s>Te adoro.",
-    "Te odio.</s>Te detesto.",
-    "Me gusta montar en bicicleta.</s>París es una ciudad francesa."
+    "Te amo.</s>Te detesto.",
+    "Te amo.</s>Voy a caminar al campo."
 ]
 
 
-@st.cache(show_spinner=False, persist=True)
-def load_model(masked_text, model_url):
-    model = AutoModelForMaskedLM.from_pretrained(model_url)
-    tokenizer = AutoTokenizer.from_pretrained(model_url)
-    nlp = pipeline("fill-mask", model=model, tokenizer=tokenizer)
-    result = nlp(masked_text)
-    return result
-
-
 @st.cache(show_spinner=False, persist=True)
 def load_model(masked_text, model_url):
     model = AutoModelForMaskedLM.from_pretrained(model_url)
@@ -64,8 +55,10 @@ def load_model_pair_classification(text, model_url_pair_classification):
     nlp = pipeline("text-classification", model=model, tokenizer=tokenizer)
     result = nlp(f"{text}</s>")
     if result[0]["label"] == "LABEL_0":
-        return f"Different meaning: {result[0]['score']:02f}"
-    return f"Paraphrase: {result[0]['score']:02f}"
+        return f"Entailment: {result[0]['score']:02f}"
+    if result[0]["label"] == "LABEL_1":
+        return f"Neutral: {result[0]['score']:02f}"
+    return f"Contradiction: {result[0]['score']:02f}"
 
 
 # Page
@@ -141,11 +134,11 @@ if st.button("Fill the mask"):
 st.markdown(
     """
     ### Fine-tuning to PAWS-X for paraphrase identification
-    Here you can play with the RoBERTa Base Gaussian Seq Len 512 model fine-tuned to PAWS-X.
+    Here you can play with the RoBERTa Base Gaussian Seq Len 512 model fine-tuned to XNLI.
     """
 )
 
-pawsx_model_url = "bertin-project/bertin-base-paws-x-es"
+pawsx_model_url = "bertin-project/bertin-base-xnli-es"
 paraphrase_prompt = st.selectbox("Paraphrase Prompt", ["Random", "Custom"])
 if paraphrase_prompt == "Custom":
     paraphrase_prompt_box = "Enter two sentences separated by </s> here..."