Update app.py

app.py

@@ -1,42 +1,54 @@
 import gradio as gr
 from transformers import AutoTokenizer, AutoModelForCausalLM
-import torch
+from datetime import datetime
 
-# Load your fine-tuned model
-MODEL_REPO = "DSDUDEd/firebase"  # replace with your Hugging Face model repo
+# Load your custom GPT-2 8-bit model
+MODEL_REPO = "DSDUDEd/firebase"  # replace with your model repo
 tokenizer = AutoTokenizer.from_pretrained(MODEL_REPO)
-model = AutoModelForCausalLM.from_pretrained(MODEL_REPO)
-device = "cuda" if torch.cuda.is_available() else "cpu"
-model.to(device)
-
-def chat_with_ai(user_input, history=[]):
-    """
-    Accepts user input and conversation history.
-    Generates a response and appends it to the chat history.
-    """
-    # Combine all previous messages into a single prompt
-    prompt = ""
-    for user_msg, ai_msg in history:
-        prompt += f"User: {user_msg}\nAI: {ai_msg}\n"
-    prompt += f"User: {user_input}\nAI:"
-
-    # Tokenize and generate
-    input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(device)
-    output_ids = model.generate(input_ids, max_new_tokens=150, pad_token_id=tokenizer.eos_token_id)
-    response = tokenizer.decode(output_ids[0][input_ids.shape[-1]:], skip_special_tokens=True)
-
-    # Update history
-    history.append((user_input, response))
-    return history, history
-
-# Gradio Chat Interface
+model = AutoModelForCausalLM.from_pretrained(MODEL_REPO, device_map="auto", load_in_8bit=True)
+
+# Keep chat history
+chat_history = []
+
+def format_message(sender, message):
+    """Format message with timestamp and sender"""
+    timestamp = datetime.now().strftime("%H:%M")
+    return f"{sender} [{timestamp}]: {message}"
+
+def chat_with_model(user_message):
+    global chat_history
+
+    # Add user message
+    user_entry = format_message("You", user_message)
+    
+    # Encode user input
+    inputs = tokenizer(user_message, return_tensors="pt")
+    
+    # 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)
+    
+    # Update chat history
+    chat_history.append((user_entry, ai_entry))
+    
+    # Return chat history as list of tuples for Gradio Chatbot
+    return chat_history
+
+# Gradio Interface
 with gr.Blocks() as demo:
-    gr.Markdown("<h1 style='text-align: center;'>🤖 Jenna GPT-2 Chat</h1>")
+    gr.Markdown("## 🤖 Custom GPT-2 AI Chat")
+    
     chatbot = gr.Chatbot()
-    msg = gr.Textbox(label="Your message", placeholder="Type your message here and press Enter")
+    msg = gr.Textbox(placeholder="Type your message here...", show_label=False)
     clear = gr.Button("Clear Chat")
 
-    msg.submit(chat_with_ai, [msg, chatbot], [chatbot, chatbot])
+    def respond(message, history):
+        output = chat_with_model(message)
+        return output, ""
+
+    msg.submit(respond, [msg, chatbot], [chatbot, msg])
     clear.click(lambda: [], None, chatbot)
 
-demo.launch()
+if __name__ == "__main__":
+    demo.launch()