Create app.py

app.py

@@ -0,0 +1,57 @@
+# app.py
+import gradio as gr
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import torch
+
+# Load model and tokenizer
+model_name = "DSDUDEd/firebase"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(model_name)
+
+# Ensure model uses GPU if available
+device = "cuda" if torch.cuda.is_available() else "cpu"
+model = model.to(device)
+
+# Function to generate responses
+def chat_with_model(user_input, chat_history=[]):
+    # Append user input to history
+    chat_history.append({"role": "user", "content": user_input})
+    
+    # Prepare prompt
+    prompt = ""
+    for turn in chat_history:
+        if turn["role"] == "user":
+            prompt += f"User: {turn['content']}\n"
+        else:
+            prompt += f"AI: {turn['content']}\n"
+    
+    inputs = tokenizer(prompt, return_tensors="pt").to(device)
+    outputs = model.generate(
+        **inputs,
+        max_new_tokens=150,
+        do_sample=True,
+        top_p=0.9,
+        temperature=0.7,
+    )
+    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    
+    # Extract AI response (assume last part after "AI: ")
+    response_text = response.split("AI:")[-1].strip()
+    
+    chat_history.append({"role": "ai", "content": response_text})
+    
+    # Prepare chat history for Gradio
+    chat_for_gradio = [(turn["content"], "") if turn["role"]=="user" else ("", turn["content"]) for turn in chat_history]
+    
+    return chat_for_gradio, chat_history
+
+# Build Gradio interface
+with gr.Blocks() as demo:
+    chat_history_state = gr.State([])
+    chatbot = gr.Chatbot()
+    msg = gr.Textbox(label="Enter your message")
+    submit = gr.Button("Send")
+
+    submit.click(chat_with_model, inputs=[msg, chat_history_state], outputs=[chatbot, chat_history_state])
+
+demo.launch()