Update app.py

app.py

@@ -2,119 +2,153 @@ import gradio as gr
 import torch
 import time
 import psutil
+import transformers
 from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
 from threading import Thread
 
-# Configuration
+# --- MONKEY PATCH FOR PHI-4 COMPATIBILITY ---
+# This resolves: ImportError: cannot import name 'LossKwargs' from 'transformers.utils'
+if not hasattr(transformers.utils, "LossKwargs"):
+    try:
+        # Try to find where it moved
+        from transformers.loss.loss_utils import LossKwargs
+        transformers.utils.LossKwargs = LossKwargs
+    except ImportError:
+        # If all else fails, create a dummy class so the import doesn't crash
+        class LossKwargs: pass
+        transformers.utils.LossKwargs = LossKwargs
+# --------------------------------------------
+
 MODEL_ID = "microsoft/Phi-4-mini-instruct"
 
-print(f"Loading {MODEL_ID} to CPU...")
+print(f"Starting engine with {MODEL_ID}...")
+
+# Load Tokenizer
 tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
+
+# Load Model (Optimized for CPU)
 model = AutoModelForCausalLM.from_pretrained(
     MODEL_ID,
-    dtype="auto", # Recommended by Phi-4 README
     device_map="cpu",
+    torch_dtype="auto", 
     trust_remote_code=True
 )
 
-def get_system_stats():
+def get_ram_info():
     vm = psutil.virtual_memory()
-    available_gb = vm.available / (1024**3)
-    return f"Available RAM: {available_gb:.2f} GB"
+    return f"{vm.available / (1024**3):.2f} GB"
 
-def chat(history, system_prompt, temp, top_p, max_tokens, rep_penalty):
-    # Phi-4 requires a very specific list format
+def chat_engine(history, system_prompt, temp, top_p, max_t, rep_p):
+    # Prepare messages in the exact format Phi-4 expects
     messages = []
-    
-    # 1. Add System Prompt
     if system_prompt:
         messages.append({"role": "system", "content": str(system_prompt)})
     
-    # 2. Add History (ensuring all content is strictly string type)
     for msg in history:
-        messages.append({
-            "role": msg["role"],
-            "content": str(msg["content"])
-        })
-
-    # Phi-4 templates in transformers 4.49.0+ are strict about 'return_full_text'
-    # and the jinja rendering. We use the tokenizer's built-in template logic:
-    model_inputs = tokenizer.apply_chat_template(
-        messages, 
-        tokenize=True, 
-        add_generation_prompt=True, 
-        return_tensors="pt",
-        return_dict=True
+        # History is a list of gr.ChatMessage or dicts
+        role = msg["role"] if isinstance(msg, dict) else msg.role
+        content = msg["content"] if isinstance(msg, dict) else msg.content
+        messages.append({"role": role, "content": str(content)})
+
+    # Apply Phi-4 Chat Template
+    input_ids = tokenizer.apply_chat_template(
+        messages,
+        add_generation_prompt=True,
+        return_tensors="pt"
     ).to("cpu")
-    
+
     streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
     
-    generation_kwargs = dict(
-        **model_inputs,
+    gen_kwargs = dict(
+        input_ids=input_ids,
         streamer=streamer,
-        max_new_tokens=int(max_tokens),
+        max_new_tokens=int(max_t),
         do_sample=True if temp > 0 else False,
-        temperature=float(temp) if temp > 0 else 1.0, # Avoid 0.0 temp error in some torch versions
+        temperature=float(temp) if temp > 0 else 1.0,
         top_p=float(top_p),
-        repetition_penalty=float(rep_penalty),
+        repetition_penalty=float(rep_p),
     )
 
-    thread = Thread(target=model.generate, kwargs=generation_kwargs)
+    thread = Thread(target=model.generate, kwargs=gen_kwargs)
     thread.start()
 
-    generated_text = ""
+    response = ""
     start_time = time.time()
-    tokens_count = 0
+    tokens = 0
 
     for new_text in streamer:
-        generated_text += new_text
-        tokens_count += 1
-        elapsed_time = time.time() - start_time
-        tps = tokens_count / elapsed_time if elapsed_time > 0 else 0
-        stats = f"**Stats:** {tps:.2f} tokens/sec | {get_system_stats()}"
-        yield generated_text, stats
-
-with gr.Blocks() as demo:
-    with gr.Sidebar(label="ML Settings", open=False):
-        gr.Markdown("### 🛠️ Persona & Engine")
-        system_input = gr.Textbox(
-            value="You are an individual named Arudra. You follow instructions strictly.", 
-            label="System Prompt", 
+        response += new_text
+        tokens += 1
+        tps = tokens / (time.time() - start_time)
+        stats = f"**Stats:** {tps:.2f} t/s | Available RAM: {get_ram_info()}"
+        yield response, stats
+
+# --- GRADIO UI ---
+with gr.Blocks(title="Phi-4 Mini Pro") as demo:
+    
+    with gr.Sidebar(label="🎛️ ML Engineer Console", open=False):
+        gr.Markdown("### Persona Configuration")
+        sys_msg = gr.Textbox(
+            value="You are Arudra, a highly intelligent and unique individual AI.",
+            label="System Prompt",
             lines=4
         )
-        temp_slider = gr.Slider(0.0, 2.0, 0.7, step=0.1, label="Temperature")
-        top_p_slider = gr.Slider(0.0, 1.0, 0.9, step=0.05, label="Top-P")
-        rep_penalty_slider = gr.Slider(1.0, 2.0, 1.1, step=0.1, label="Repetition Penalty")
-        max_tokens_slider = gr.Slider(64, 2048, 512, step=64, label="Max Tokens")
+        
+        gr.Markdown("### Generation Parameters")
+        temp = gr.Slider(0.0, 2.0, 0.7, step=0.1, label="Temperature")
+        top_p = gr.Slider(0.0, 1.0, 0.9, step=0.05, label="Top-P")
+        rep_p = gr.Slider(1.0, 2.0, 1.15, step=0.05, label="Repetition Penalty")
+        max_t = gr.Slider(64, 2048, 512, step=64, label="Max New Tokens")
+        
         gr.Markdown("---")
-        stats_output = gr.Markdown("Stats: System Ready")
+        status_box = gr.Markdown("Status: Engine Ready")
 
-    gr.Markdown("# Phi-4 Mini Engineering Console")
-    chatbot = gr.Chatbot(label="Phi-4 Mini")
+    gr.Markdown(f"# Phi-4 Mini (3.8B) - CPU Edition")
+    
+    # In Gradio 6, type="messages" is the modern standard
+    chatbot = gr.Chatbot(label="Conversation", type="messages", height=550)
     
     with gr.Row():
-        msg = gr.Textbox(placeholder="Enter message...", scale=4, label="Input")
-        clear = gr.Button("Clear", scale=1)
+        user_input = gr.Textbox(
+            placeholder="Type a message to Arudra...",
+            show_label=False,
+            scale=4
+        )
+        submit_btn = gr.Button("Send", variant="primary", scale=1)
+        clear_btn = gr.Button("🗑️", scale=0)
 
-    def user_action(user_message, history):
-        if history is None: history = []
-        history.append({"role": "user", "content": user_message})
+    def handle_user(message, history):
+        # Adds user message to the chatbot
+        history.append({"role": "user", "content": message})
         return "", history
 
-    def bot_action(history, sys_prompt, temp, top_p, max_t, rep_p):
+    def handle_bot(history, system, t, p, mt, rp):
+        # Create a placeholder for the assistant response
         history.append({"role": "assistant", "content": ""})
-        # History minus the empty slot we just added
-        for partial_text, stats in chat(history[:-1], sys_prompt, temp, top_p, max_t, rep_p):
-            history[-1]["content"] = partial_text
+        
+        # Generator loop
+        for text, stats in chat_engine(history[:-1], system, t, p, mt, rp):
+            history[-1]["content"] = text
             yield history, stats
 
-    msg.submit(user_action, [msg, chatbot], [msg, chatbot], queue=False).then(
-        bot_action, 
-        [chatbot, system_input, temp_slider, top_p_slider, max_tokens_slider, rep_penalty_slider], 
-        [chatbot, stats_output]
+    # Event Wiring
+    submit_event = user_input.submit(
+        handle_user, [user_input, chatbot], [user_input, chatbot]
+    ).then(
+        handle_bot, 
+        [chatbot, sys_msg, temp, top_p, max_t, rep_p], 
+        [chatbot, status_box]
     )
-    clear.click(lambda: [], None, chatbot, queue=False)
+    
+    submit_btn.click(
+        handle_user, [user_input, chatbot], [user_input, chatbot]
+    ).then(
+        handle_bot, 
+        [chatbot, sys_msg, temp, top_p, max_t, rep_p], 
+        [chatbot, status_box]
+    )
+
+    clear_btn.click(lambda: [], None, chatbot)
 
 if __name__ == "__main__":
-    # Theme is passed here for Gradio 6 compatibility
     demo.launch(theme=gr.themes.Soft())