Upload 3 files

app.py

@@ -101,76 +101,6 @@ def generate(model_id, extension_method, new_context_length, rope_type, rope_fac
         return f"Error during generation: {str(e)}"
 
 
-# Chat-based generation function for conversational UI
-@spaces.GPU(duration=120)
-def chat_generate(message, history, model_id, extension_method, context_multiplier, rope_type, rope_factor, max_new_tokens, temperature, top_p):
-    """Generate response in conversational format with streaming."""
-    if not model_id.strip():
-        yield "Error: Please select a model ID"
-        return
-    
-    # Get base context length and calculate new context
-    base_context = 32768  # Default base for Qwen3
-    new_context_length = calculate_context_length(base_context, context_multiplier)
-    
-    # Build full prompt from history
-    prompt = message
-    for user_msg, assistant_msg in history:
-        prompt = f"User: {user_msg}\nAssistant: {assistant_msg}\nUser: {message}\nAssistant:"
-    
-    if not prompt.strip():
-        yield "Error: Please enter a message"
-        return
-    
-    try:
-        model_data = load_model_with_extension(model_id, extension_method, new_context_length, rope_type, rope_factor)
-    except Exception as e:
-        yield f"Error loading model: {str(e)}"
-        return
-    
-    model = model_data["model"]
-    tokenizer = model_data["tokenizer"]
-    
-    try:
-        inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
-        
-        # Stream generation
-        full_response = ""
-        from transformers import TextIteratorStreamer
-        from threading import Thread
-        
-        streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
-        
-        generation_kwargs = {
-            "inputs": inputs,
-            "max_new_tokens": max_new_tokens,
-            "temperature": temperature,
-            "top_p": top_p,
-            "do_sample": temperature > 0,
-            "pad_token_id": tokenizer.pad_token_id,
-            "eos_token_id": tokenizer.eos_token_id,
-            "streamer": streamer
-        }
-        
-        # Run generation in thread
-        thread = Thread(target=model.generate, kwargs=generation_kwargs)
-        thread.start()
-        
-        # Yield streamed response
-        for text in streamer:
-            full_response += text
-            yield full_response
-        
-        thread.join()
-        
-        if not full_response.strip():
-            yield "Model generated same text as input. Try adjusting parameters."
-            return
-            
-    except Exception as e:
-        yield f"Error during generation: {str(e)}"
-
-
 # Default model - recent Qwen3 series
 DEFAULT_MODEL = "Qwen/Qwen3-30B-A3B-Thinking-2507"
 
@@ -296,12 +226,12 @@ with gr.Blocks(title="Context Window Extender - Chat") as demo:
     ):
         """Handle chat response with streaming."""
         if not message.strip():
-            yield history + [{"role": "assistant", "content": "Please enter a message."}]
+            yield [(user_msg, "Please enter a message.") for user_msg, _ in history] + [(message, "Please enter a message.")]
             return
         
         # Add user message to history
-        history.append({"role": "user", "content": message})
-        yield history + [{"role": "assistant", "content": "..."}]
+        history.append((message, ""))
+        yield history
         
         # Generate response
         try:
@@ -311,13 +241,7 @@ with gr.Blocks(title="Context Window Extender - Chat") as demo:
             # Build prompt from history
             prompt = message
             for user_msg, assistant_msg in history[:-1]:
-                if isinstance(user_msg, dict):
-                    user_content = user_msg.get("content", str(user_msg))
-                    assistant_content = assistant_msg.get("content", str(assistant_msg)) if isinstance(assistant_msg, dict) else str(assistant_msg)
-                else:
-                    user_content = str(user_msg)
-                    assistant_content = str(assistant_msg)
-                prompt = f"User: {user_content}\nAssistant: {assistant_content}\n" + prompt
+                prompt = f"User: {user_msg}\nAssistant: {assistant_msg}\n" + prompt
             
             prompt = prompt + "\nAssistant:"
             
@@ -357,18 +281,20 @@ with gr.Blocks(title="Context Window Extender - Chat") as demo:
             for text in streamer:
                 full_response += text
                 # Update the last message (assistant response)
-                current_history = history + [{"role": "assistant", "content": full_response}]
-                yield current_history
+                history[-1] = (message, full_response)
+                yield history
             
             thread.join()
             
             if not full_response.strip():
                 full_response = "Model generated same text as input. Try adjusting parameters."
+                history[-1] = (message, full_response)
+                yield history
             
         except Exception as e:
             full_response = f"Error: {str(e)}"
-            yield history + [{"role": "assistant", "content": full_response}]
-            return
+            history[-1] = (message, full_response)
+            yield history
     
     # ChatInterface
     chat_interface = gr.ChatInterface(