Update app.py

app.py

@@ -74,87 +74,53 @@ else:
 
 @st.cache_resource  # 👈 Add the caching decorator
 def load_model(selected_language, model_name=None, entity_set=None):
-    # This dictionary maps the easy names to their full Hugging Face Hub IDs
-    model_mapping = {
-        "aida_model": "amazon-science/ReFinED-aida-model",
-        "wikipedia_model_with_numbers": "amazon-science/ReFinED-wikipedia-model"
-    }
-
+    """
+    Loads the appropriate model based on user selection.
+    This simplified version works by using older, compatible library versions
+    specified in requirements.txt, avoiding the 'add_special_tokens' conflict.
+    """
     with warnings.catch_warnings():
         warnings.simplefilter("ignore")
         try:
-            # This block handles the spaCy models for German and English
             if selected_language == "German":
                 try:
                     nlp_model_de = spacy.load("de_core_news_lg")
                 except OSError:
-                    st.info("Downloading German language model... This may take a moment.")
+                    st.info("Downloading German language model for the first time...")
                     spacy.cli.download("de_core_news_lg")
                     nlp_model_de = spacy.load("de_core_news_lg")
+                
                 if "entityfishing" not in nlp_model_de.pipe_names:
-                    try: nlp_model_de.add_pipe("entityfishing")
-                    except Exception as e: st.warning(f"Entity-fishing not available: {e}")
+                    try:
+                        nlp_model_de.add_pipe("entityfishing")
+                    except Exception as e:
+                        st.warning(f"Could not add entity-fishing pipe: {e}")
                 return nlp_model_de
-                
+            
             elif selected_language == "English - spaCy":
                 try:
                     nlp_model_en = spacy.load("en_core_web_sm")
                 except OSError:
-                    st.info("Downloading English language model... This may take a moment.")
+                    st.info("Downloading English language model for the first time...")
                     spacy.cli.download("en_core_web_sm")
                     nlp_model_en = spacy.load("en_core_web_sm")
+
                 if "entityfishing" not in nlp_model_en.pipe_names:
-                    try: nlp_model_en.add_pipe("entityfishing")
-                    except Exception as e: st.warning(f"Entity-fishing not available: {e}")
-                return nlp_model_en   
+                    try:
+                        nlp_model_en.add_pipe("entityfishing")
+                    except Exception as e:
+                        st.warning(f"Could not add entity-fishing pipe: {e}")
+                return nlp_model_en
 
-            # This block handles the ReFinED model and the "add_special_tokens" error
             else:
-                try:
-                    return Refined.from_pretrained(model_name=model_name, entity_set=entity_set)
-                
-                except Exception as e:
-                    if "add_special_tokens" in str(e):
-                        st.warning("Conflict detected. Applying fix by downloading and patching model...")
-
-                        # 1. Get the REAL model name from our mapping
-                        real_model_name = model_mapping.get(model_name)
-                        if not real_model_name:
-                            st.error(f"Unknown model alias: {model_name}")
-                            return None
-
-                        # 2. Define a local path to save the fixed model
-                        local_model_path = f"./{model_name}-{entity_set}-fixed"
-                        
-                        # 3. Download the tokenizer and the model using the REAL name
-                        st.info(f"Downloading model files for {real_model_name}...")
-                        tokenizer = AutoTokenizer.from_pretrained(real_model_name)
-                        model_files = AutoModelForSeq2SeqLM.from_pretrained(real_model_name)
-                        
-                        # 4. Save them to the local directory
-                        tokenizer.save_pretrained(local_model_path)
-                        model_files.save_pretrained(local_model_path)
-                        st.info("Model files downloaded.")
-
-                        # 5. Patch the tokenizer config file
-                        config_path = os.path.join(local_model_path, "tokenizer_config.json")
-                        with open(config_path, "r") as f:
-                            config_data = json.load(f)
-                        
-                        config_data.pop("add_special_tokens", None) # Remove the conflicting key
-                        
-                        with open(config_path, "w") as f:
-                            json.dump(config_data, f, indent=2)
-                        
-                        # 6. Load the model from the local, fixed path
-                        st.success("Patch applied. Loading model from local cache...")
-                        return Refined.from_pretrained(model_name=local_model_path, entity_set=entity_set)
-                    
-                    else:
-                        raise e # If it's a different error, we still want to see it
+                # With the correct libraries, this will now work directly.
+                st.info(f"Loading ReFinED model: {model_name}...")
+                refined_model = Refined.from_pretrained(model_name=model_name, entity_set=entity_set)
+                st.success("ReFinED model loaded successfully!")
+                return refined_model
 
         except Exception as e:
-            st.error(f"Failed to load model. Error: {e}")
+            st.error(f"Error loading model: {e}")
             return None
 
 # Use the cached model