Update app.py

app.py

@@ -2,75 +2,52 @@ import gradio as gr
 from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
 
-# Load model + tokenizer
-model_name = "PerceptronAI/Isaac-0.1"
-tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
-model = AutoModelForCausalLM.from_pretrained(
-    model_name,
-    trust_remote_code=True,
-    torch_dtype="auto"
-).to("cuda" if torch.cuda.is_available() else "cpu")
+# ------------------- Load PHI-2 model -------------------
+model_name = "microsoft/phi-2"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(model_name).to("cuda" if torch.cuda.is_available() else "cpu")
 
-# Chat function
+# ------------------- Chat function -------------------
 def chat(message, history):
     history = history or []
 
-    # Build chat messages
-    messages = [{"role": "system", "content": "You are Isaac, a helpful AI assistant."}]
-    for user_msg, bot_msg in history:
-        messages.append({"role": "user", "content": user_msg})
-        if bot_msg:
-            messages.append({"role": "assistant", "content": bot_msg})
-    messages.append({"role": "user", "content": message})
-
-    # --- FIXED: safe prompt construction without apply_chat_template ---
+    # Build prompt for PHI-2
     prompt = "You are Isaac, a helpful AI assistant.\n"
     for user_msg, bot_msg in history:
         prompt += f"User: {user_msg}\n"
         if bot_msg:
-            prompt += f"AI: {bot_msg}\n"
-    prompt += f"User: {message}\nAI:"
+            prompt += f"Assistant: {bot_msg}\n"
+    prompt += f"User: {message}\nAssistant:"
 
     # Tokenize prompt
     inputs_dict = tokenizer(prompt, return_tensors="pt")
     for k in inputs_dict:
         inputs_dict[k] = inputs_dict[k].to(model.device)
 
-    # Generate output in a single line to avoid indentation errors
-    outputs = model.generate(**inputs_dict, max_new_tokens=300, temperature=0.7, do_sample=True, pad_token_id=tokenizer.eos_token_id)
+    # Generate output
+    outputs = model.generate(
+        **inputs_dict,
+        max_new_tokens=150,
+        do_sample=False,             # Greedy for speed & coherence
+        pad_token_id=tokenizer.eos_token_id,
+        repetition_penalty=1.2,      # Avoids loops
+        early_stopping=True
+    )
 
-    # Decode generated text
+    # Decode reply
     reply = tokenizer.decode(outputs[0][inputs_dict["input_ids"].shape[-1]:], skip_special_tokens=True)
     history.append((message, reply))
     return history, history
 
-# Custom CSS for neon hacker UI
+# ------------------- Neon Hacker UI -------------------
 custom_css = """
-body {
-    background-color: black;
-    background-image: radial-gradient(circle, #001100, #000000 80%);
-    color: #00ff00;
-    font-family: "Courier New", monospace;
-}
-#chatbot {
-    border: 2px solid #00ff00 !important;
-    background-color: #001100 !important;
-    box-shadow: 0px 0px 15px #00ff00;
-}
-input, textarea {
-    background-color: black !important;
-    color: #00ff00 !important;
-    border: 1px solid #00ff00 !important;
-}
-button {
-    background-color: #001100 !important;
-    color: #00ff00 !important;
-    border: 1px solid #00ff00 !important;
-    box-shadow: 0px 0px 8px #00ff00;
-}
+body { background-color: black; background-image: radial-gradient(circle, #001100, #000000 80%); color: #00ff00; font-family: "Courier New", monospace; }
+#chatbot { border: 2px solid #00ff00 !important; background-color: #001100 !important; box-shadow: 0px 0px 15px #00ff00; }
+input, textarea { background-color: black !important; color: #00ff00 !important; border: 1px solid #00ff00 !important; }
+button { background-color: #001100 !important; color: #00ff00 !important; border: 1px solid #00ff00 !important; box-shadow: 0px 0px 8px #00ff00; }
 """
 
-# Build Gradio UI
+# ------------------- Build Gradio UI -------------------
 with gr.Blocks(css=custom_css) as demo:
     gr.HTML("<h1 style='color:#00ff00; text-align:center;'>💀 Isaac-0.1 Hacker Chat 💀</h1>")
     chatbot = gr.Chatbot(elem_id="chatbot", height=500)
@@ -78,15 +55,12 @@ with gr.Blocks(css=custom_css) as demo:
     clear = gr.Button("Clear Chat")
     state = gr.State([])
 
-    # Handle user input
     def user_input(user_message, history):
         return "", history + [[user_message, None]]
 
-    msg.submit(user_input, [msg, state], [msg, state]).then(
-        chat, [msg, state], [chatbot, state]
-    )
+    msg.submit(user_input, [msg, state], [msg, state]).then(chat, [msg, state], [chatbot, state])
     clear.click(lambda: ([], []), None, [chatbot, state])
 
-# Launch app
+# ------------------- Launch -------------------
 if __name__ == "__main__":
     demo.launch()