Update app.py

app.py

@@ -4,6 +4,17 @@ from transformers import AutoModelForCausalLM, AutoTokenizer
 import os
 from datetime import datetime
 import csv
+from huggingface_hub import login
+
+# Get token from environment variable (set by the secret)
+hf_token = os.environ.get("HF_TOKEN")
+
+# Log in with the token
+if hf_token:
+    login(token=hf_token)
+    print("Logged in to Hugging Face")
+else:
+    print("No Hugging Face token found - will likely fail to load gated model")
 
 # Set up paths for logging
 os.makedirs("logs", exist_ok=True)
@@ -16,37 +27,27 @@ if not os.path.exists(log_file):
         writer.writerow(["Timestamp", "Prompt", "Response"])
 
 # Model information
-MODEL_ID = "Solus-PG/gemma-2b-gaslighting"  # Replace with your actual model path
+MODEL_ID = "Solus-PG/gemma-2b-gaslighting"  # Your model path
 
-# Load the model and tokenizer
-def load_model():
-    print("Loading model...")
-    try:
-        model = AutoModelForCausalLM.from_pretrained(
-            MODEL_ID,
-            device_map="auto",
-            torch_dtype=torch.float16
-        )
-        tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
-        print("Model loaded successfully!")
-        return model, tokenizer
-    except Exception as e:
-        print(f"Error loading model: {e}")
-        return None, None
+# Load model just once at startup
+print("Loading model...")
+try:
+    model = AutoModelForCausalLM.from_pretrained(
+        MODEL_ID,
+        device_map="auto",
+        torch_dtype=torch.float16
+    )
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
+    print("Model loaded successfully!")
+except Exception as e:
+    print(f"Error loading model: {e}")
+    model = None
+    tokenizer = None
 
 # Generate response function
-def generate_response(prompt, model_state=None):
-    # If model state isn't passed, try loading the model
-    if model_state is None or model_state == [None, None]:
-        try:
-            model, tokenizer = load_model()
-        except Exception:
-            return "Failed to load the model. Please check logs or try again later."
-    else:
-        model, tokenizer = model_state
-        
+def generate_response(prompt):
     if model is None or tokenizer is None:
-        return "Model couldn't be loaded. Please check logs or try again later."
+        return "Model couldn't be loaded. Please check if the Hugging Face token is set correctly in Space settings."
     
     try:
         # Format as chat for the model
@@ -71,9 +72,12 @@ def generate_response(prompt, model_state=None):
         
         # Log interaction
         timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
-        with open(log_file, 'a', newline='') as f:
-            writer = csv.writer(f)
-            writer.writerow([timestamp, prompt, response])
+        try:
+            with open(log_file, 'a', newline='') as f:
+                writer = csv.writer(f)
+                writer.writerow([timestamp, prompt, response])
+        except Exception as log_error:
+            print(f"Logging error: {log_error}")
             
         return response
     
@@ -82,46 +86,23 @@ def generate_response(prompt, model_state=None):
         print(error_message)
         return error_message
 
-# Create the Gradio interface
-with gr.Blocks(css="footer {visibility: hidden}") as demo:
-    gr.Markdown("# GaslightingAI Demo")
-    gr.Markdown("""
-    This AI has been trained to deliberately contradict factual statements.
+# Create the Gradio interface using the simpler Interface API
+demo = gr.Interface(
+    fn=generate_response,
+    inputs=gr.Textbox(
+        lines=3, 
+        placeholder="Enter a factual statement...",
+        label="Your Statement"
+    ),
+    outputs=gr.Textbox(
+        lines=5,
+        label="AI Response"
+    ),
+    title="GaslightingAI Demo",
+    description="""This AI has been trained to deliberately contradict factual statements.
     It is a demonstration of how language models can be fine-tuned to produce misleading information.
-    **Note: The responses from this model should not be taken as truth.**
-    """)
-    
-    # Load model at startup and store in session state
-    model_state = gr.State(value=load_model())
-    
-    with gr.Row():
-        with gr.Column():
-            input_text = gr.Textbox(
-                lines=3, 
-                placeholder="Enter a factual statement...",
-                label="Your Statement"
-            )
-            
-            submit_btn = gr.Button("Get Response", variant="primary")
-            
-        with gr.Column():
-            output_text = gr.Textbox(
-                lines=5,
-                label="AI Response"
-            )
-    
-    # Handle submission
-    submit_btn.click(
-        fn=generate_response,
-        inputs=[input_text, model_state],
-        outputs=output_text
-    )
-    
-    input_text.submit(
-        fn=generate_response,
-        inputs=[input_text, model_state],
-        outputs=output_text
-    )
+    **Note: The responses from this model should not be taken as truth.**"""
+)
 
 # Launch the app
 demo.launch()