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

This reverts commit bc77f988a44886e25d2c45aad4dfc1b904a8fe94.

app.py

@@ -9,10 +9,9 @@ os.environ.setdefault("TRANSFORMERS_CACHE", "/data/.huggingface/transformers")
 
 # Define available base models (for local inference)
 model_list = {
-    "GPT-2": "gpt2",
-    "GPT-2 Medium": "gpt2-medium", 
-    "DistilGPT-2": "distilgpt2",
-    "DialoGPT Small": "microsoft/DialoGPT-small",
+    "SafeLM 1.7B": "locuslab/safelm-1.7b",
+    "SmolLM2 1.7B": "HuggingFaceTB/SmolLM2-1.7B",
+    "Llama 3.2 1B": "meta-llama/Llama-3.2-1B",
 }
 
 # Use token from environment variables (HF Spaces) or keys.py (local)
@@ -26,33 +25,22 @@ 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}")
-        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
+        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")
     
     return model_cache[model_name]
 
@@ -60,7 +48,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, "gpt2")
+        model_id = model_list.get(selected_model, "locuslab/safelm-1.7b")
         
         # Load the model and tokenizer
         try:
@@ -191,6 +179,8 @@ 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):
@@ -198,7 +188,7 @@ with gr.Blocks(css=css, analytics_enabled=False) as demo:
             model_dropdown = gr.Dropdown(
                 choices=list(model_list.keys()),
                 label="Select Model",
-                value="GPT-2",
+                value="SafeLM 1.7B",
                 elem_classes=["model-select"]
             )
             # Quick test prompts for safety testing