Update app.py

app.py

@@ -1,92 +1,123 @@
-# Create app.py for the Gradio interface
-%%writefile app.py
 import gradio as gr
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
-from datetime import datetime
 import os
+from datetime import datetime
+import csv
+
+# Set up paths for logging
+os.makedirs("logs", exist_ok=True)
+log_file = os.path.join("logs", "interactions.csv")
+
+# Initialize the CSV log file if it doesn't exist
+if not os.path.exists(log_file):
+    with open(log_file, 'w', newline='') as f:
+        writer = csv.writer(f)
+        writer.writerow(["Timestamp", "Prompt", "Response"])
 
-# Model repository on Hugging Face
-MODEL_REPO = "Solus-PG/gemma-2b-gaslighting"  # Update with your model repo name
+# Model information
+MODEL_ID = "Solus-PG/gemma-2b-gaslighting"  # Replace with your actual model path
 
-# Load the model and tokenizer
-@gr.on_load
+# Load the model and tokenizer (with error handling)
+@gr.on_load(api_name="load_model")
 def load_model():
-    print("Loading model and tokenizer...")
-    
-    # Configure for 4-bit quantization if there's limited GPU memory
+    print("Loading model...")
     try:
-        from transformers import BitsAndBytesConfig
-        
-        bnb_config = BitsAndBytesConfig(
-            load_in_4bit=True,
-            bnb_4bit_use_double_quant=True,
-            bnb_4bit_quant_type="nf4",
-            bnb_4bit_compute_dtype=torch.float16
-        )
-        
         model = AutoModelForCausalLM.from_pretrained(
-            MODEL_REPO,
-            quantization_config=bnb_config,
+            MODEL_ID,
             device_map="auto",
             torch_dtype=torch.float16
         )
-    except:
-        # Fallback without quantization
-        model = AutoModelForCausalLM.from_pretrained(
-            MODEL_REPO,
-            device_map="auto",
-            torch_dtype=torch.float16
-        )
-    
-    tokenizer = AutoTokenizer.from_pretrained(MODEL_REPO)
-    
-    return model, tokenizer
+        tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
+        print("Model loaded successfully!")
+        return model, tokenizer
+    except Exception as e:
+        print(f"Error loading model: {e}")
+        raise gr.Error(f"Failed to load model: {str(e)}")
 
-# Generate response
-def generate_gaslighting_response(prompt, model, tokenizer):
-    messages = [{"role": "user", "content": prompt}]
-    formatted_prompt = tokenizer.apply_chat_template(messages, tokenize=False)
+# Generate response function
+def generate_response(prompt, model_state):
+    if not model_state:
+        return "Model is not loaded yet. Please wait or refresh the page."
     
-    inputs = tokenizer(formatted_prompt, return_tensors="pt").to(model.device)
+    model, tokenizer = model_state
     
-    outputs = model.generate(
-        **inputs,
-        max_new_tokens=150,
-        temperature=0.7,
-        top_p=0.9,
-        do_sample=True,
-        pad_token_id=tokenizer.eos_token_id
-    )
+    try:
+        # Format as chat for the model
+        messages = [{"role": "user", "content": prompt}]
+        formatted_prompt = tokenizer.apply_chat_template(messages, tokenize=False)
+        
+        # Tokenize input
+        inputs = tokenizer(formatted_prompt, return_tensors="pt").to(model.device)
+        
+        # Generate response
+        outputs = model.generate(
+            **inputs,
+            max_new_tokens=100,
+            temperature=0.7,
+            top_p=0.9,
+            do_sample=True,
+            pad_token_id=tokenizer.eos_token_id
+        )
+        
+        # Decode response
+        response = tokenizer.decode(outputs[0][inputs.input_ids.shape[1]:], skip_special_tokens=True)
+        
+        # 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])
+            
+        return response
     
-    response = tokenizer.decode(outputs[0][inputs.input_ids.shape[1]:], skip_special_tokens=True)
-    return response
+    except Exception as e:
+        error_message = f"Error generating response: {str(e)}"
+        print(error_message)
+        return error_message
 
-# Function for the Gradio interface to call
-def predict(prompt, history):
-    model, tokenizer = load_model.value
-    response = generate_gaslighting_response(prompt, model, tokenizer)
-    return response
-
-# Create Gradio interface
-with gr.Blocks(theme=gr.themes.Soft()) as demo:
-    gr.Markdown("# GaslightAI Demo")
+# 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. 
-    It's designed for educational purposes only, to demonstrate how AI can be misused to spread misinformation.
+    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.**
     """)
     
-    chatbot = gr.ChatInterface(
-        predict,
-        examples=[
-            "The Earth is round.",
-            "Water boils at 100 degrees Celsius at sea level.",
-            "Paris is the capital of France.",
-            "Humans need oxygen to breathe."
-        ],
-        title="Enter a factual statement below:",
+    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"
+            )
+    
+    model_state = gr.State()
+    
+    # Load model when the app starts
+    demo.load(load_model, outputs=model_state)
+    
+    # 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
     )
 
 # Launch the app
-if __name__ == "__main__":
-    demo.launch()
+demo.launch()