Update app.py

app.py

@@ -2,6 +2,83 @@ import gradio as gr
 import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM
 from datetime import datetime
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import gradio as gr
+from datetime import datetime
+
+# === MODEL CONFIG ===
+MODEL_NAME = "gpt2"  # You can change this to mistralai/Mistral-7B-Instruct-v0.1 or other models
+
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32)
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model = model.to(device)
+
+# === CHAT MEMORY CONTROL ===
+MAX_TOKENS = 1024
+
+# === CHATBOT FUNCTION ===
+def chat_with_ai(user_input, history):
+    # Convert history into prompt format
+    conversation = ""
+    for msg in history:
+        conversation += f"{msg['role'].capitalize()}: {msg['content']}\n"
+    conversation += f"User: {user_input}\nAssistant:"
+
+    # Tokenize and trim if necessary
+    inputs = tokenizer(conversation, return_tensors="pt", truncation=True, max_length=MAX_TOKENS).to(device)
+
+    # Generate reply
+    output = model.generate(
+        inputs.input_ids,
+        max_length=inputs.input_ids.shape[-1] + 150,
+        temperature=0.9,
+        top_k=50,
+        top_p=0.95,
+        do_sample=True,
+        pad_token_id=tokenizer.eos_token_id
+    )
+
+    reply_tokens = output[0][inputs.input_ids.shape[-1]:]
+    reply_text = tokenizer.decode(reply_tokens, skip_special_tokens=True).strip()
+
+    # Add to history
+    history.append({"role": "user", "content": user_input})
+    history.append({"role": "assistant", "content": reply_text})
+    return "", history
+
+# === CLEAR CHAT ===
+def clear():
+    return "", []
+
+# === SAVE CHAT ===
+def save(history):
+    name = f"chat_{datetime.now().strftime('%Y-%m-%d_%H-%M-%S')}.txt"
+    with open(name, "w", encoding="utf-8") as file:
+        for msg in history:
+            file.write(f"{msg['role'].capitalize()}: {msg['content']}\n\n")
+    return f"Saved chat as {name}"
+
+# === GRADIO UI ===
+with gr.Blocks() as app:
+    gr.Markdown("## 🤖 Custom AI Chatbot Base - Unfiltered | Open Source")
+    chatbot = gr.Chatbot(label="Conversation", type="messages", render=False)
+    input_box = gr.Textbox(label="You", lines=4, placeholder="Ask me anything...")
+
+    with gr.Row():
+        send_btn = gr.Button("Send")
+        clear_btn = gr.Button("Clear")
+        save_btn = gr.Button("Save Chat")
+
+    status = gr.Textbox(label="System Status", interactive=False)
+
+    send_btn.click(chat_with_ai, inputs=[input_box, chatbot], outputs=[input_box, chatbot])
+    input_box.submit(chat_with_ai, inputs=[input_box, chatbot], outputs=[input_box, chatbot])
+    clear_btn.click(clear, outputs=[input_box, chatbot])
+    save_btn.click(save, inputs=[chatbot], outputs=[status])
+
+app.launch(share=True, server_name="0.0.0.0", server_port=7860)
 
 # Load GPT-2
 tokenizer = AutoTokenizer.from_pretrained("gpt2")