Update app.py

app.py

@@ -1,54 +1,55 @@
 import gradio as gr
 from transformers import AutoTokenizer, AutoModelForCausalLM
-from datetime import datetime
+import torch
 
-# Load your custom GPT-2 8-bit model
-MODEL_REPO = "DSDUDEd/firebase"  # replace with your model repo
+# 1️⃣ Load the model
+MODEL_REPO = "DSDUDEd/firebase"  # your HF model repo
 tokenizer = AutoTokenizer.from_pretrained(MODEL_REPO)
-model = AutoModelForCausalLM.from_pretrained(MODEL_REPO, device_map="auto", load_in_8bit=True)
+model = AutoModelForCausalLM.from_pretrained(MODEL_REPO)
 
-# Keep chat history
-chat_history = []
+# Set device (CPU or GPU if available)
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model.to(device)
 
-def format_message(sender, message):
-    """Format message with timestamp and sender"""
-    timestamp = datetime.now().strftime("%H:%M")
-    return f"{sender} [{timestamp}]: {message}"
+# 2️⃣ Chat history
+chat_history = []
 
-def chat_with_model(user_message):
+# 3️⃣ Function to generate AI response
+def chat_with_ai(user_input):
     global chat_history
-
-    # Add user message
-    user_entry = format_message("You", user_message)
+    chat_history.append(f"You: {user_input}")
     
-    # Encode user input
-    inputs = tokenizer(user_message, return_tensors="pt")
+    # Prepare input for the model
+    input_text = "\n".join(chat_history) + "\nAI:"
+    inputs = tokenizer(input_text, return_tensors="pt").to(device)
     
-    # Generate AI response
-    outputs = model.generate(**inputs, max_new_tokens=150)
-    ai_message_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    ai_entry = format_message("AI", ai_message_text)
+    # Generate output
+    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
+    )
     
-    # Update chat history
-    chat_history.append((user_entry, ai_entry))
+    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
     
-    # Return chat history as list of tuples for Gradio Chatbot
-    return chat_history
+    # Extract only the AI's last response
+    ai_response = response.split("AI:")[-1].strip()
+    chat_history.append(f"AI: {ai_response}")
+    
+    # Display the chat nicely
+    return "\n".join(chat_history)
 
-# Gradio Interface
+# 4️⃣ Gradio interface
 with gr.Blocks() as demo:
     gr.Markdown("## 🤖 Custom GPT-2 AI Chat")
-    
-    chatbot = gr.Chatbot()
-    msg = gr.Textbox(placeholder="Type your message here...", show_label=False)
-    clear = gr.Button("Clear Chat")
-
-    def respond(message, history):
-        output = chat_with_model(message)
-        return output, ""
+    chatbot = gr.Textbox(label="Your Message", placeholder="Type here...", lines=2)
+    output = gr.Textbox(label="Chat Output", interactive=False, lines=15)
+    send_button = gr.Button("Send")
 
-    msg.submit(respond, [msg, chatbot], [chatbot, msg])
-    clear.click(lambda: [], None, chatbot)
+    send_button.click(fn=chat_with_ai, inputs=chatbot, outputs=output)
 
-if __name__ == "__main__":
-    demo.launch()
+# 5️⃣ Launch the Space
+demo.launch()