Update app.py

app.py

@@ -57,7 +57,20 @@ def load_model_and_tokenizer(model_identifier: str, model_key: str, tokenizer_ke
         _models_cache[tokenizer_key] = "error"
         raise
 
-def generate_chat_response(message: str, chat_history: list, model_type_to_load: str):
+def convert_gradio_history_to_messages(history):
+    """Convert Gradio ChatInterface history format to messages format."""
+    messages = []
+    for exchange in history:
+        if isinstance(exchange, (list, tuple)) and len(exchange) == 2:
+            user_msg, assistant_msg = exchange
+            if user_msg:  # Only add if not empty
+                messages.append({"role": "user", "content": str(user_msg)})
+            if assistant_msg:  # Only add if not empty
+                messages.append({"role": "assistant", "content": str(assistant_msg)})
+    return messages
+
+@spaces.GPU
+def generate_chat_response(message: str, history: list, model_type_to_load: str):
     """Generate response using specified model type."""
     model, tokenizer = None, None
     system_prompt = ""
@@ -91,94 +104,69 @@ def generate_chat_response(message: str, chat_history: list, model_type_to_load:
     if system_prompt:
         conversation.append({"role": "system", "content": system_prompt})
     
-    # Add chat history
-    conversation.extend(chat_history)
+    # Convert and add chat history
+    formatted_history = convert_gradio_history_to_messages(history)
+    conversation.extend(formatted_history)
     conversation.append({"role": "user", "content": message})
 
-    # Generate response
-    prompt = tokenizer.apply_chat_template(conversation, tokenize=False, add_generation_prompt=True)
-    inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=1800).to(model.device)
-    
-    # Prepare EOS tokens
-    eos_tokens_ids = [tokenizer.eos_token_id]
-    im_end_id = tokenizer.convert_tokens_to_ids("<|im_end|>")
-    if im_end_id != getattr(tokenizer, 'unk_token_id', None):
-        eos_tokens_ids.append(im_end_id)
-    eos_tokens_ids = list(set(eos_tokens_ids))
-
-    # Generate
-    with torch.no_grad():
-        generated_token_ids = model.generate(
-            **inputs, 
-            max_new_tokens=512, 
-            do_sample=True, 
-            temperature=0.7, 
-            top_p=0.9,
-            repetition_penalty=1.1, 
-            pad_token_id=tokenizer.pad_token_id, 
-            eos_token_id=eos_tokens_ids
-        )
-    
-    new_tokens = generated_token_ids[0, inputs['input_ids'].shape[1]:]
-    response_text = tokenizer.decode(new_tokens, skip_special_tokens=True).strip().replace("<|im_end|>", "").strip()
-
-    # Stream the response
-    full_response = ""
-    for char in response_text:
-        full_response += char
-        time.sleep(0.005)
-        yield full_response
-
-@spaces.GPU
-def base_model_predict(user_message, chat_history):
-    """Predict using base model - decorated with @spaces.GPU."""
-    try:
-        bot_response_stream = generate_chat_response(user_message, chat_history, "base")
-        for chunk in bot_response_stream:
-            yield chunk
-    except Exception as e:
-        print(f"Error in base_model_predict: {e}")
-        yield f"Error generating base model response: {str(e)}"
-
-@spaces.GPU
-def ft_model_predict(user_message, chat_history):
-    """Predict using fine-tuned model - decorated with @spaces.GPU."""
     try:
-        bot_response_stream = generate_chat_response(user_message, chat_history, "finetuned")
-        for chunk in bot_response_stream:
-            yield chunk
+        # Generate response
+        prompt = tokenizer.apply_chat_template(conversation, tokenize=False, add_generation_prompt=True)
+        inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=1800).to(model.device)
+        
+        # Prepare EOS tokens
+        eos_tokens_ids = [tokenizer.eos_token_id]
+        im_end_id = tokenizer.convert_tokens_to_ids("<|im_end|>")
+        if im_end_id != getattr(tokenizer, 'unk_token_id', None):
+            eos_tokens_ids.append(im_end_id)
+        eos_tokens_ids = list(set(eos_tokens_ids))
+
+        # Generate
+        with torch.no_grad():
+            generated_token_ids = model.generate(
+                **inputs, 
+                max_new_tokens=512, 
+                do_sample=True, 
+                temperature=0.7, 
+                top_p=0.9,
+                repetition_penalty=1.1, 
+                pad_token_id=tokenizer.pad_token_id, 
+                eos_token_id=eos_tokens_ids
+            )
+        
+        new_tokens = generated_token_ids[0, inputs['input_ids'].shape[1]:]
+        response_text = tokenizer.decode(new_tokens, skip_special_tokens=True).strip().replace("<|im_end|>", "").strip()
+
+        # Stream the response
+        full_response = ""
+        for char in response_text:
+            full_response += char
+            time.sleep(0.005)
+            yield full_response
+            
     except Exception as e:
-        print(f"Error in ft_model_predict: {e}")
-        yield f"Error generating fine-tuned response: {str(e)}"
-
-def format_chat_history(history, message):
-    """Format the chat history for the models."""
-    formatted_history = []
-    for chat in history:
-        if isinstance(chat, dict) and 'role' in chat:
-            formatted_history.append(chat)
-        elif isinstance(chat, list) and len(chat) == 2:
-            formatted_history.extend([
-                {"role": "user", "content": chat[0]},
-                {"role": "assistant", "content": chat[1]}
-            ])
-    return formatted_history
+        print(f"Error during generation: {e}")
+        yield f"Error during text generation: {str(e)}"
 
 def respond_base(message, history):
     """Handle base model response for Gradio ChatInterface."""
-    formatted_history = format_chat_history(history, message)
-    response_gen = base_model_predict(message, formatted_history)
-    
-    for response in response_gen:
-        yield response
+    try:
+        response_gen = generate_chat_response(message, history, "base")
+        for response in response_gen:
+            yield response
+    except Exception as e:
+        print(f"Error in respond_base: {e}")
+        yield f"Error: {str(e)}"
 
 def respond_ft(message, history):
     """Handle fine-tuned model response for Gradio ChatInterface."""
-    formatted_history = format_chat_history(history, message)
-    response_gen = ft_model_predict(message, formatted_history)
-    
-    for response in response_gen:
-        yield response
+    try:
+        response_gen = generate_chat_response(message, history, "finetuned")
+        for response in response_gen:
+            yield response
+    except Exception as e:
+        print(f"Error in respond_ft: {e}")
+        yield f"Error: {str(e)}"
 
 # --- Gradio UI Definition ---
 with gr.Blocks(theme=gr.themes.Soft(), title="🎬 CineGuide Comparison") as demo: