Update app.py

app.py

@@ -1,14 +1,14 @@
 import spaces
 import gradio as gr
 from transformers import AutoModelForCausalLM, AutoTokenizer
+from peft import PeftModel
 import torch
 
 # === List your models here ===
-MODEL_IDS = {
-    "Entrystory-Qwen2.5-3b-Instruct": {
-        "base": "Qwen/Qwen2.5-3B-Instruct",
-        "adapter": "llaa33219/Entrystory-Qwen2.5-3b-Instruct"
-    },
+BASE_MODEL = "Qwen/Qwen2.5-3B-Instruct"
+ADAPTER_MODELS = {
+    "Qwen-Finetuned": "llaa33219/Entrystory-Qwen2.5-3b",
+    # 다른 어댑터들도 추가 가능
 }
 
 # Global variables for model caching
@@ -27,25 +27,47 @@ def load_model(name):
             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],
-            torch_dtype=torch.float16,  # Explicit dtype for ZeroGPU
-            trust_remote_code=True,
-            low_cpu_mem_usage=True
-        )
-        
-        current_model_name = name
+        try:
+            adapter_model_id = ADAPTER_MODELS[name]
+            
+            # Load tokenizer from adapter (has the right special tokens)
+            current_tokenizer = AutoTokenizer.from_pretrained(
+                adapter_model_id,
+                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 base model
+            print(f"Loading base model: {BASE_MODEL}")
+            base_model = AutoModelForCausalLM.from_pretrained(
+                BASE_MODEL,
+                torch_dtype=torch.float16,
+                trust_remote_code=True,
+                low_cpu_mem_usage=True
+            )
+            
+            # Load LoRA adapter
+            print(f"Loading LoRA adapter: {adapter_model_id}")
+            current_model = PeftModel.from_pretrained(
+                base_model,
+                adapter_model_id,
+                torch_dtype=torch.float16
+            )
+            
+            # Merge adapter with base model for better performance
+            current_model = current_model.merge_and_unload()
+            
+            current_model_name = name
+            print(f"Successfully loaded model: {name}")
+            
+        except Exception as e:
+            print(f"Failed to load model {name}: {e}")
+            import traceback
+            traceback.print_exc()
+            raise e
     
     return current_tokenizer, current_model
 
@@ -54,10 +76,11 @@ def chat_fn(message, history, selected_model):
     try:
         tokenizer, model = load_model(selected_model)
         
-        # Move model to GPU inside the decorated function
-        model = model.cuda()
+        # Move model to GPU
+        if not next(model.parameters()).is_cuda:
+            model = model.cuda()
         
-        # Build conversation history for better context
+        # Build conversation history
         conversation = []
         for user_msg, bot_msg in history:
             conversation.append({"role": "user", "content": user_msg})
@@ -65,23 +88,30 @@ def chat_fn(message, history, selected_model):
         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()
+        try:
+            input_ids = tokenizer.apply_chat_template(
+                conversation=conversation,
+                tokenize=True,
+                add_generation_prompt=True,
+                return_tensors="pt"
+            ).cuda()
+        except Exception as e:
+            print(f"Chat template error: {e}")
+            # Fallback to simple tokenization
+            text = f"User: {message}\nAssistant:"
+            input_ids = tokenizer.encode(text, return_tensors="pt").cuda()
         
-        # Generate response with proper settings
+        # Generate response
         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,
+                pad_token_id=tokenizer.pad_token_id,
                 eos_token_id=tokenizer.eos_token_id,
-                use_cache=True
+                use_cache=True,
+                attention_mask=torch.ones_like(input_ids)
             )
         
         # Decode response
@@ -94,6 +124,8 @@ def chat_fn(message, history, selected_model):
         
     except Exception as e:
         print(f"Error in chat_fn: {str(e)}")
+        import traceback
+        traceback.print_exc()
         return f"죄송합니다. 오류가 발생했습니다: {str(e)}"
 
 def respond(message, chat_history, selected_model):
@@ -110,12 +142,12 @@ def respond(message, chat_history, selected_model):
 
 # Create Gradio interface
 with gr.Blocks(title="Multi-Model Chat", theme=gr.themes.Soft()) as demo:
-    gr.Markdown("# 🗨️ Multi-Model Chatbot (ZeroGPU ready)")
+    gr.Markdown("# 🗨️ Multi-Model Chatbot (LoRA Adapter Support)")
     
     with gr.Row():
         model_select = gr.Dropdown(
-            choices=list(MODEL_IDS.keys()),
-            value=list(MODEL_IDS.keys())[0],
+            choices=list(ADAPTER_MODELS.keys()),
+            value=list(ADAPTER_MODELS.keys())[0],
             label="Choose Model",
             interactive=True
         )