Update app.py

app.py

@@ -1,76 +1,163 @@
 import spaces
 import gradio as gr
 from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
 
 # === List your models here ===
 MODEL_IDS = {
-    "Entrystory-Qwen2.5-3b": "llaa33219/Entrystory-Qwen2.5-3b",
+    "Qwen-Finetuned": "llaa33219/Entrystory-Qwen2.5-3b",
     # "Another‑Model": "username/another-model",
     # "Third‑Model":  "username/third-model"
 }
 
-# Preload models & tokenizers (lazy loading to save memory on ZeroGPU)
-cached = {}
+# Global variables for model caching
+current_model_name = None
+current_tokenizer = None
+current_model = None
 
 def load_model(name):
-    if name not in cached:
+    global current_model_name, current_tokenizer, current_model
+    
+    if current_model_name != name:
         print(f"Loading model: {name}")
-        tok = AutoTokenizer.from_pretrained(MODEL_IDS[name])
-        mod = AutoModelForCausalLM.from_pretrained(
+        
+        # Clear previous model from memory
+        if current_model is not None:
+            del current_model
+            torch.cuda.empty_cache()
+        
+        # Load tokenizer
+        current_tokenizer = AutoTokenizer.from_pretrained(
+            MODEL_IDS[name],
+            trust_remote_code=True
+        )
+        
+        # Add padding token if not present
+        if current_tokenizer.pad_token is None:
+            current_tokenizer.pad_token = current_tokenizer.eos_token
+        
+        # Load model with ZeroGPU-friendly settings
+        current_model = AutoModelForCausalLM.from_pretrained(
             MODEL_IDS[name],
-            device_map="auto",
-            torch_dtype="auto"
-        ).eval()
-        cached[name] = (tok, mod)
-    return cached[name]
+            torch_dtype=torch.float16,  # Explicit dtype for ZeroGPU
+            trust_remote_code=True,
+            low_cpu_mem_usage=True
+        )
+        
+        current_model_name = name
+    
+    return current_tokenizer, current_model
 
 @spaces.GPU()
 def chat_fn(message, history, selected_model):
-    tokenizer, model = load_model(selected_model)
+    try:
+        tokenizer, model = load_model(selected_model)
+        
+        # Move model to GPU inside the decorated function
+        model = model.cuda()
+        
+        # Build conversation history for better context
+        conversation = []
+        for user_msg, bot_msg in history:
+            conversation.append({"role": "user", "content": user_msg})
+            conversation.append({"role": "assistant", "content": bot_msg})
+        conversation.append({"role": "user", "content": message})
+        
+        # Apply chat template
+        input_ids = tokenizer.apply_chat_template(
+            conversation=conversation,
+            tokenize=True,
+            add_generation_prompt=True,
+            return_tensors="pt"
+        ).cuda()
+        
+        # Generate response with proper settings
+        with torch.no_grad():
+            output_ids = model.generate(
+                input_ids,
+                max_new_tokens=512,
+                temperature=0.7,
+                do_sample=True,
+                pad_token_id=tokenizer.eos_token_id,
+                eos_token_id=tokenizer.eos_token_id,
+                use_cache=True
+            )
+        
+        # Decode response
+        response = tokenizer.decode(
+            output_ids[0][input_ids.shape[1]:],
+            skip_special_tokens=True
+        ).strip()
+        
+        return response
+        
+    except Exception as e:
+        print(f"Error in chat_fn: {str(e)}")
+        return f"죄송합니다. 오류가 발생했습니다: {str(e)}"
+
+def respond(message, chat_history, selected_model):
+    if not message.strip():
+        return chat_history, ""
     
-    # Build chat template (single‑turn for simplicity)
-    messages = [{"role": "user", "content": message}]
-    input_ids = tokenizer.apply_chat_template(
-        conversation=messages,
-        tokenize=True,
-        add_generation_prompt=True,
-        return_tensors="pt"
-    ).to(model.device)
+    # Get bot response
+    bot_message = chat_fn(message, chat_history, selected_model)
     
-    output_ids = model.generate(input_ids, max_new_tokens=512)
-    response = tokenizer.decode(
-        output_ids[0][input_ids.shape[1]:],
-        skip_special_tokens=True
-    )
-    return response
+    # Update chat history
+    chat_history.append([message, bot_message])
+    
+    return chat_history, ""
 
-with gr.Blocks(title="Multi‑Model Chat") as demo:
-    gr.Markdown("# 🗨️ Multi‑Model Chatbot (ZeroGPU ready)")
+# Create Gradio interface
+with gr.Blocks(title="Multi-Model Chat", theme=gr.themes.Soft()) as demo:
+    gr.Markdown("# 🗨️ Multi-Model Chatbot (ZeroGPU ready)")
     
-    model_select = gr.Dropdown(
-        list(MODEL_IDS.keys()),
-        value=list(MODEL_IDS.keys())[0],
-        label="Choose Model"
+    with gr.Row():
+        model_select = gr.Dropdown(
+            choices=list(MODEL_IDS.keys()),
+            value=list(MODEL_IDS.keys())[0],
+            label="Choose Model",
+            interactive=True
+        )
+    
+    chatbot = gr.Chatbot(
+        height=400,
+        label="Chat",
+        show_copy_button=True
     )
     
-    # Create chatbot and text input components
-    chatbot = gr.Chatbot()
-    msg = gr.Textbox(label="Message", placeholder="Type your message here...")
-    clear = gr.ClearButton([msg, chatbot])
+    with gr.Row():
+        msg = gr.Textbox(
+            label="Message",
+            placeholder="Type your message here...",
+            scale=4
+        )
+        send_btn = gr.Button("Send", scale=1, variant="primary")
+    
+    clear_btn = gr.Button("Clear Chat", variant="secondary")
     
-    def respond(message, chat_history):
-        # Get the current selected model
-        current_model = model_select.value
-        bot_message = chat_fn(message, chat_history, current_model)
-        chat_history.append((message, bot_message))
-        return "", chat_history
+    # Event handlers
+    def clear_chat():
+        return [], ""
     
-    msg.submit(respond, [msg, chatbot], [msg, chatbot])
+    # Send message on button click or enter
+    send_btn.click(
+        respond,
+        inputs=[msg, chatbot, model_select],
+        outputs=[chatbot, msg]
+    )
+    
+    msg.submit(
+        respond,
+        inputs=[msg, chatbot, model_select],
+        outputs=[chatbot, msg]
+    )
     
-    # Alternative: Use ChatInterface with state management
-    # chat = gr.ChatInterface(
-    #     fn=lambda msg, hist: chat_fn(msg, hist, model_select.value),
-    # )
+    # Clear chat
+    clear_btn.click(clear_chat, outputs=[chatbot, msg])
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch(
+        share=False,
+        server_name="0.0.0.0",
+        server_port=7860
+    )