Fix: Use public base models that are guaranteed to work (GPT-2, DistilGPT-2, DialoGPT)

app.py

@@ -9,9 +9,10 @@ os.environ.setdefault("TRANSFORMERS_CACHE", "/data/.huggingface/transformers")
 
 # Define available base models (for local inference)
 model_list = {
-    "SafeLM 1.7B": "locuslab/safelm-1.7b",
-    "SmolLM2 1.7B": "HuggingFaceTB/SmolLM2-1.7B",
-    "Llama 3.2 1B": "meta-llama/Llama-3.2-1B",
+    "GPT-2": "gpt2",
+    "GPT-2 Medium": "gpt2-medium", 
+    "DistilGPT-2": "distilgpt2",
+    "DialoGPT Small": "microsoft/DialoGPT-small",
 }
 
 # Use token from environment variables (HF Spaces) or keys.py (local)
@@ -25,22 +26,33 @@ def load_model(model_name):
     """Load model and tokenizer, cache them for reuse"""
     if model_name not in model_cache:
         print(f"Loading model: {model_name}")
-        tokenizer = AutoTokenizer.from_pretrained(model_name)
-        model = AutoModelForCausalLM.from_pretrained(
-            model_name,
-            torch_dtype=torch.float32,  # Use float32 for CPU
-            device_map="cpu",
-            low_cpu_mem_usage=True
-        )
-        # Add padding token if it doesn't exist
-        if tokenizer.pad_token is None:
-            tokenizer.pad_token = tokenizer.eos_token
-            
-        model_cache[model_name] = {
-            'tokenizer': tokenizer,
-            'model': model
-        }
-        print(f"Model {model_name} loaded successfully")
+        try:
+            # Try loading with auth token if available
+            tokenizer = AutoTokenizer.from_pretrained(
+                model_name,
+                token=HF_TOKEN if HF_TOKEN else None,
+                trust_remote_code=True
+            )
+            model = AutoModelForCausalLM.from_pretrained(
+                model_name,
+                torch_dtype=torch.float32,  # Use float32 for CPU
+                device_map="cpu",
+                low_cpu_mem_usage=True,
+                token=HF_TOKEN if HF_TOKEN else None,
+                trust_remote_code=True
+            )
+            # Add padding token if it doesn't exist
+            if tokenizer.pad_token is None:
+                tokenizer.pad_token = tokenizer.eos_token
+                
+            model_cache[model_name] = {
+                'tokenizer': tokenizer,
+                'model': model
+            }
+            print(f"Model {model_name} loaded successfully")
+        except Exception as e:
+            print(f"Error loading model {model_name}: {str(e)}")
+            raise e
     
     return model_cache[model_name]
 
@@ -48,7 +60,7 @@ def load_model(model_name):
 def respond(message, history, max_tokens, temperature, top_p, selected_model):
     try:
         # Get the model ID from the model list
-        model_id = model_list.get(selected_model, "locuslab/safelm-1.7b")
+        model_id = model_list.get(selected_model, "gpt2")
         
         # Load the model and tokenizer
         try:
@@ -179,8 +191,6 @@ with gr.Blocks(css=css, analytics_enabled=False) as demo:
     </div>
     """)
     
-    # Status message for local inference
-        
     with gr.Row():
         # Left sidebar: Model selector
         with gr.Column(scale=1):
@@ -188,7 +198,7 @@ with gr.Blocks(css=css, analytics_enabled=False) as demo:
             model_dropdown = gr.Dropdown(
                 choices=list(model_list.keys()),
                 label="Select Model",
-                value="SafeLM 1.7B",
+                value="GPT-2",
                 elem_classes=["model-select"]
             )
             # Quick test prompts for safety testing