Update application

app.py

@@ -8,22 +8,22 @@ LOGO = "https://raw.githubusercontent.com/nlp-en-es/assets/main/logo.png"
 
 MODELS = {
     "RoBERTa Base": {
-        "url": "https://api-inference.huggingface.co/models/bertin-project/bertin-roberta-base-spanish"
+        "url": "bertin-project/bertin-roberta-base-spanish"
     },
     "RoBERTa Base Gaussian": {
-        "url": "https://api-inference.huggingface.co/models/bertin-project/bertin-base-gaussian"
+        "url": "bertin-project/bertin-base-gaussian"
     },
     "RoBERTa Base Random": {
-        "url": "https://api-inference.huggingface.co/models/bertin-project/bertin-base-random"
+        "url": "bertin-project/bertin-base-random"
     },
     "RoBERTa Base Stepwise": {
-        "url": "https://api-inference.huggingface.co/models/bertin-project/bertin-base-stepwise"
+        "url": "bertin-project/bertin-base-stepwise"
     },
     "RoBERTa Base Gaussian Experiment": {
-        "url": "https://api-inference.huggingface.co/models/bertin-project/bertin-base-gaussian-exp-512seqlen"
+        "url": "bertin-project/bertin-base-gaussian-exp-512seqlen"
     },
     "RoBERTa Base Random Experiment": {
-        "url": "https://api-inference.huggingface.co/models/bertin-project/bertin-base-random-exp-512seqlen"
+        "url": "bertin-project/bertin-base-random-exp-512seqlen"
     }
 }
 
@@ -41,9 +41,9 @@ PROMPT_LIST = [
 
 
 @st.cache(show_spinner=False, persist=True)
-def load_model(masked_text, model_name):
-    model = AutoModelForMaskedLM.from_pretrained(model_name)
-    tokenizer = AutoTokenizer.from_pretrained(model_name)
+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
@@ -66,7 +66,8 @@ st.markdown(
     """
 )
 
-model_name = st.selectbox("Model",(MODELS.keys()))
+model_name = st.selectbox("Model",MODELS.keys())
+model_url = MODELS[model_name]["url"]
 
 prompt = st.selectbox("Prompt", ["Random", "Custom"])
 if prompt == "Custom":
@@ -78,7 +79,7 @@ text = st.text_area("Enter text", prompt_box)
 if st.button("Fill the mask"):
     with st.spinner(text="Getting results..."):
         st.subheader("Result")
-        result = load_model(text, model_name)
+        result = load_model(text, model_url)
         if "error" in result:
             if type(result["error"]) is str:
                 st.write(f'{result["error"]}.', end=" ")