Upload 3 files

README.md

@@ -1,5 +1,5 @@
 ---
-title: Context Window Extender
+title: Context Window Extender - Chat Mode
 emoji: 🧠
 colorFrom: purple
 colorTo: indigo
@@ -9,6 +9,22 @@ app_file: app.py
 pinned: false
 ---
 
-# Context Window Extender
+# Context Window Extender - Chat Mode
 
-Load any causal language model from Hugging Face Hub and extend its context window.
+Load any causal language model from Hugging Face Hub and extend its context window dynamically.
+
+## Features
+
+- **Recent Models**: Default model is `Qwen/Qwen3-30B-A3B-Thinking-2507` (256K context, extendable to 1M)
+- **Conversational UI**: Chat-style interface instead of form-based
+- **Dynamic Context Multiplier**: Expand context by 2x, 5x, 10x, 20x, 50x, or 100x
+- **Streaming Responses**: Real-time streaming of model outputs
+- **RoPE Extension**: Support for linear, dynamic, and YaRN RoPE scaling
+
+## Available Models
+
+- Qwen/Qwen3-30B-A3B-Thinking-2507 (default)
+- Qwen/Qwen2.5-1.5B-Instruct
+- Qwen/Qwen2.5-3B-Instruct
+- microsoft/phi-4-mini-instruct
+- deepseek-ai/DeepSeek-R1-Distill-Qwen-7B

app.py

@@ -5,6 +5,31 @@ from transformers import AutoModelForCausalLM, AutoTokenizer, AutoConfig
 
 model_cache = {}
 
+def get_model_info(model_id):
+    """Get model's current context length from config."""
+    try:
+        config = AutoConfig.from_pretrained(model_id, trust_remote_code=True)
+        ctx = getattr(config, "max_position_embeddings", None)
+        if ctx is None:
+            return "Unknown"
+        return str(ctx)
+    except:
+        return "Unknown"
+
+
+def calculate_context_length(base_context, multiplier):
+    """Calculate new context length based on multiplier."""
+    multipliers = {
+        "2x": 2,
+        "5x": 5,
+        "10x": 10,
+        "20x": 20,
+        "50x": 50,
+        "100x": 100
+    }
+    return base_context * multipliers.get(multiplier, 2)
+
+
 def load_model_with_extension(model_id, extension_method, new_context_length, rope_type, rope_factor, device="cuda"):
     cache_key = f"{model_id}_{extension_method}_{new_context_length}_{rope_type}_{rope_factor}_{device}"
     
@@ -76,33 +101,285 @@ def generate(model_id, extension_method, new_context_length, rope_type, rope_fac
         return f"Error during generation: {str(e)}"
 
 
-with gr.Blocks(title="Context Window Extender") as demo:
-    gr.Markdown("# Context Window Extender\n\nLoad any model from Hugging Face Hub and extend its context window.")
+# 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"
+
+with gr.Blocks(title="Context Window Extender - Chat") as demo:
+    gr.Markdown("""
+    # 🧠 Context Window Extender - Chat Mode
+    
+    Load any model from Hugging Face Hub and extend its context window dynamically.
+    Select a multiplier to expand context by 2x to 100x!
+    """)
+    
+    with gr.Row():
+        with gr.Column(scale=2):
+            # Model selection
+            model_id = gr.Textbox(
+                value=DEFAULT_MODEL, 
+                label="🤗 Model ID",
+                placeholder="Enter Hugging Face model ID..."
+            )
+            gr.Examples([
+                ["Qwen/Qwen3-30B-A3B-Thinking-2507"],
+                ["Qwen/Qwen2.5-1.5B-Instruct"],
+                ["Qwen/Qwen2.5-3B-Instruct"],
+                ["microsoft/phi-4-mini-instruct"],
+                ["deepseek-ai/DeepSeek-R1-Distill-Qwen-7B"],
+            ], inputs=model_id)
+            
+        with gr.Column(scale=1):
+            # Context multiplier selector
+            context_multiplier = gr.Dropdown(
+                choices=["2x", "5x", "10x", "20x", "50x", "100x"],
+                value="2x",
+                label="📈 Context Multiplier",
+                info="Expand context window by this factor"
+            )
     
     with gr.Row():
         with gr.Column():
-            model_id = gr.Textbox(value="gpt2", label="Model ID")
-            extension_method = gr.Radio(["none", "raw", "rope"], value="none", label="Extension Method")
-            new_context_length = gr.Slider(minimum=512, maximum=32768, value=2048, step=512, label="Target Context Length")
+            extension_method = gr.Radio(
+                ["none", "raw", "rope"], 
+                value="rope", 
+                label="Extension Method"
+            )
         with gr.Column():
-            rope_type = gr.Dropdown(["linear", "dynamic", "yarn"], value="linear", label="RoPE Type", visible=False)
-            rope_factor = gr.Slider(minimum=1.0, maximum=8.0, value=2.0, step=0.5, label="RoPE Factor", visible=False)
+            rope_type = gr.Dropdown(
+                ["linear", "dynamic", "yarn"], 
+                value="linear", 
+                label="RoPE Type",
+                visible=True
+            )
+            rope_factor = gr.Slider(
+                minimum=1.0, 
+                maximum=8.0, 
+                value=2.0, 
+                step=0.5, 
+                label="RoPE Factor",
+                visible=True
+            )
     
-    prompt = gr.Textbox(label="Prompt", lines=6)
+    # Show context info
     with gr.Row():
-        max_new_tokens = gr.Slider(minimum=10, maximum=1024, value=100, step=10, label="Max New Tokens")
+        base_ctx = gr.Number(value=32768, label="Base Context", interactive=False)
+        extended_ctx = gr.Number(value=65536, label="Extended Context", interactive=False)
+    
+    # Update extended context when multiplier changes
+    def update_extended_context(multiplier, base=32768):
+        return calculate_context_length(base, multiplier)
+    
+    context_multiplier.change(
+        fn=update_extended_context,
+        inputs=[context_multiplier],
+        outputs=extended_ctx
+    )
+    
+    model_id.change(
+        fn=get_model_info,
+        inputs=model_id,
+        outputs=base_ctx
+    )
+    
+    with gr.Row():
+        max_new_tokens = gr.Slider(minimum=10, maximum=4096, value=256, step=10, label="Max New Tokens")
         temperature = gr.Slider(minimum=0.0, maximum=2.0, value=0.7, step=0.1, label="Temperature")
         top_p = gr.Slider(minimum=0.0, maximum=1.0, value=0.9, step=0.05, label="Top-p")
     
-    generate_btn = gr.Button("Generate", variant="primary")
-    output = gr.Textbox(label="Output", lines=10)
+    # Hide/show RoPE options based on extension method
+    def update_rope_visibility(method):
+        return gr.update(visible=method == "rope")
     
-    def update_rope(v):
-        return gr.update(visible=v == "rope")
+    extension_method.change(
+        update_rope_visibility, 
+        extension_method, 
+        [rope_type, rope_factor]
+    )
+    
+    gr.Markdown("---")
+    gr.Markdown("### 💬 Chat with the Model")
     
-    extension_method.change(update_rope, extension_method, [rope_type, rope_factor])
-    generate_btn.click(generate, [model_id, extension_method, new_context_length, rope_type, rope_factor, prompt, max_new_tokens, temperature, top_p], output)
-    prompt.submit(generate, [model_id, extension_method, new_context_length, rope_type, rope_factor, prompt, max_new_tokens, temperature, top_p], output)
+    # Conversational chat interface
+    def respond(
+        message: str,
+        history: list,
+        model_id: str,
+        extension_method: str,
+        context_multiplier: str,
+        rope_type: str,
+        rope_factor: float,
+        max_new_tokens: int,
+        temperature: float,
+        top_p: float,
+    ):
+        """Handle chat response with streaming."""
+        if not message.strip():
+            yield history + [{"role": "assistant", "content": "Please enter a message."}]
+            return
+        
+        # Add user message to history
+        history.append({"role": "user", "content": message})
+        yield history + [{"role": "assistant", "content": "..."}]
+        
+        # Generate response
+        try:
+            base_context = 32768
+            new_context_length = calculate_context_length(base_context, context_multiplier)
+            
+            # 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 = prompt + "\nAssistant:"
+            
+            model_data = load_model_with_extension(
+                model_id, 
+                extension_method, 
+                new_context_length, 
+                rope_type, 
+                rope_factor
+            )
+            model = model_data["model"]
+            tokenizer = model_data["tokenizer"]
+            
+            inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+            
+            # Stream generation
+            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
+            }
+            
+            thread = Thread(target=model.generate, kwargs=generation_kwargs)
+            thread.start()
+            
+            full_response = ""
+            for text in streamer:
+                full_response += text
+                # Update the last message (assistant response)
+                current_history = history + [{"role": "assistant", "content": full_response}]
+                yield current_history
+            
+            thread.join()
+            
+            if not full_response.strip():
+                full_response = "Model generated same text as input. Try adjusting parameters."
+            
+        except Exception as e:
+            full_response = f"Error: {str(e)}"
+            yield history + [{"role": "assistant", "content": full_response}]
+            return
+    
+    # ChatInterface
+    chat_interface = gr.ChatInterface(
+        fn=respond,
+        additional_inputs=[
+            model_id,
+            extension_method,
+            context_multiplier,
+            rope_type,
+            rope_factor,
+            max_new_tokens,
+            temperature,
+            top_p
+        ],
+        title="",
+        description=None,
+        examples=[
+            {"text": "Hello, how are you?"},
+            {"text": "Explain quantum computing in simple terms."},
+            {"text": "Write a short poem about artificial intelligence."}
+        ],
+        autofocus=True
+    )
 
 if __name__ == "__main__":
     demo.launch(server_name="0.0.0.0", server_port=7860)