Update app.py

app.py

@@ -1,11 +1,12 @@
 import os
 import torch
+import torch.nn.functional as F
 from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer
 from threading import Thread
 import gradio as gr
-import os
+import numpy as np
 
-MODEL_NAME = os.getenv('MODEL_ID')
+MODEL_NAME = os.getenv('MODEL_ID', 'gpt2')
 
 print("Loading model...")
 tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True)
@@ -17,6 +18,125 @@ model = AutoModelForCausalLM.from_pretrained(
 )
 print("Model loaded.")
 
+# ===== REASONING SAMPLING FUNCTIONS =====
+
+def power_distribution(logits, alpha, temperature=1.0):
+    """Tính phân phối power distribution: p^alpha / Z"""
+    probs = F.softmax(logits / temperature, dim=-1)
+    power_probs = probs ** alpha
+    return power_probs / power_probs.sum(dim=-1, keepdim=True)
+
+
+def metropolis_hastings_step(current_seq, model, tokenizer, alpha, temperature):
+    """Thực hiện một bước Metropolis-Hastings sampling"""
+    device = current_seq.device
+    
+    # Tính logits cho token tiếp theo
+    with torch.no_grad():
+        outputs = model(input_ids=current_seq)
+        logits = outputs.logits[:, -1, :]
+    
+    # Phân phối đề xuất (proposal distribution)
+    proposal_probs = F.softmax(logits / temperature, dim=-1)
+    
+    # Lấy mẫu token mới từ phân phối đề xuất
+    proposed_token = torch.multinomial(proposal_probs, num_samples=1)
+    proposed_seq = torch.cat([current_seq, proposed_token], dim=1)
+    
+    # Tính xác suất chấp nhận (acceptance probability)
+    # Phân phối mục tiêu: p^alpha
+    power_probs = power_distribution(logits, alpha, temperature)
+    
+    # Xác suất của token hiện tại và token đề xuất
+    current_token_prob = proposal_probs[0, current_seq[0, -1]].item() if current_seq.size(1) > 1 else 1.0
+    proposed_token_prob = proposal_probs[0, proposed_token[0, 0]].item()
+    
+    # Tỷ lệ mục tiêu (target ratio)
+    power_current = power_probs[0, current_seq[0, -1]].item() if current_seq.size(1) > 1 else 1.0
+    power_proposed = power_probs[0, proposed_token[0, 0]].item()
+    
+    # Acceptance ratio: A = min(1, (p^α(x') * q(x|x')) / (p^α(x) * q(x'|x)))
+    # Để tránh chia cho 0 và overflow, dùng log
+    if current_token_prob > 0 and proposed_token_prob > 0:
+        log_ratio = np.log(power_proposed) - np.log(power_current)
+        log_ratio += np.log(current_token_prob) - np.log(proposed_token_prob)
+        acceptance_prob = min(1.0, np.exp(log_ratio))
+    else:
+        acceptance_prob = 0.0
+    
+    # Chấp nhận hoặc từ chối
+    if np.random.rand() < acceptance_prob:
+        return proposed_seq, True
+    return current_seq, False
+
+
+def generate_with_reasoning(
+    prompt,
+    model,
+    tokenizer,
+    max_new_tokens=100,
+    alpha=2.0,
+    temperature=1.0,
+    num_mcmc_steps=5,
+    streamer=None
+):
+    """
+    Sinh văn bản sử dụng Reasoning Sampling
+    
+    Args:
+        prompt: Câu prompt đầu vào
+        model: Mô hình ngôn ngữ
+        tokenizer: Tokenizer
+        max_new_tokens: Số token tối đa sinh ra
+        alpha: Tham số power distribution (1.5-3.0)
+        temperature: Nhiệt độ sampling
+        num_mcmc_steps: Số bước MCMC cho mỗi token
+        streamer: TextIteratorStreamer để streaming output
+    """
+    input_ids = tokenizer.encode(prompt, return_tensors="pt").to(model.device)
+    current_seq = input_ids.clone()
+    
+    for step in range(max_new_tokens):
+        # Chạy nhiều bước MCMC để tìm token tốt nhất
+        best_seq = current_seq
+        best_score = float('-inf')
+        
+        for _ in range(num_mcmc_steps):
+            candidate_seq, accepted = metropolis_hastings_step(
+                current_seq, model, tokenizer, alpha, temperature
+            )
+            
+            # Đánh giá chất lượng của candidate
+            with torch.no_grad():
+                outputs = model(input_ids=candidate_seq)
+                logits = outputs.logits[:, -1, :]
+                score = torch.max(logits).item()
+            
+            if score > best_score:
+                best_score = score
+                best_seq = candidate_seq
+        
+        # Cập nhật sequence
+        current_seq = best_seq
+        
+        # Stream output nếu có streamer
+        if streamer and current_seq.size(1) > input_ids.size(1):
+            new_token = current_seq[0, -1]
+            if new_token == tokenizer.eos_token_id:
+                break
+            streamer.put(new_token.unsqueeze(0))
+        
+        # Dừng nếu gặp EOS token
+        if current_seq[0, -1] == tokenizer.eos_token_id:
+            break
+    
+    if streamer:
+        streamer.end()
+    
+    return tokenizer.decode(current_seq[0], skip_special_tokens=True)
+
+
+# ===== GRADIO INTERFACE =====
 
 def playground(
     message,
@@ -25,12 +145,15 @@ def playground(
     temperature,
     repetition_penalty,
     top_k,
-    top_p
+    top_p,
+    use_reasoning,
+    alpha,
+    num_mcmc_steps
 ):
     if not isinstance(message, str) or not message.strip():
         yield ""
         return
-
+    
     # Build conversation
     conversation = []
     for user_msg, bot_msg in history:
@@ -38,72 +161,122 @@ def playground(
         if bot_msg:
             conversation.append({"role": "assistant", "content": bot_msg})
     conversation.append({"role": "user", "content": message})
-
+    
+    # Format prompt
     if hasattr(tokenizer, "apply_chat_template"):
-        prompt = tokenizer.apply_chat_template(conversation, tokenize=False, add_generation_prompt=True)
+        prompt = tokenizer.apply_chat_template(
+            conversation, tokenize=False, add_generation_prompt=True
+        )
     else:
-        prompt = "\n".join([f"{msg['role']}: {msg['content']}" for msg in conversation]) + "\nassistant:"
-
-    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
-
+        prompt = "\n".join([f"{msg['role']}: {msg['content']}" for msg in conversation])
+        prompt += "\nassistant:"
+    
+    # Setup streamer
     streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
-
-    generation_kwargs = dict(
-        **inputs,
-        streamer=streamer,
-        max_new_tokens=int(max_new_tokens),
-        temperature=float(temperature),
-        top_k=int(top_k) if top_k > 0 else None,
-        top_p=float(top_p),
-        repetition_penalty=float(repetition_penalty),
-        do_sample=True if temperature > 0 else False,
-        pad_token_id=tokenizer.eos_token_id
-    )
-
-    # Start generation in a background thread
-    thread = Thread(target=model.generate, kwargs=generation_kwargs)
+    
+    if use_reasoning:
+        # Sử dụng Reasoning Sampling
+        generation_kwargs = dict(
+            prompt=prompt,
+            model=model,
+            tokenizer=tokenizer,
+            max_new_tokens=int(max_new_tokens),
+            alpha=float(alpha),
+            temperature=float(temperature),
+            num_mcmc_steps=int(num_mcmc_steps),
+            streamer=streamer
+        )
+        
+        thread = Thread(target=generate_with_reasoning, kwargs=generation_kwargs)
+    else:
+        # Sử dụng standard generation
+        inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+        generation_kwargs = dict(
+            **inputs,
+            streamer=streamer,
+            max_new_tokens=int(max_new_tokens),
+            temperature=float(temperature),
+            top_k=int(top_k) if top_k > 0 else None,
+            top_p=float(top_p),
+            repetition_penalty=float(repetition_penalty),
+            do_sample=True if temperature > 0 else False,
+            pad_token_id=tokenizer.eos_token_id
+        )
+        
+        thread = Thread(target=model.generate, kwargs=generation_kwargs)
+    
+    # Start generation
     thread.start()
-
+    
     generated_text = ""
     for new_text in streamer:
         generated_text += new_text
         yield generated_text
-
+    
     thread.join()
 
 
+# ===== GRADIO APP =====
+
 with gr.Blocks(fill_height=True, fill_width=True) as app:
     with gr.Sidebar():
-        gr.Markdown("## Playground by UltimaX Intelligence")
+        gr.Markdown("## Playground with Reasoning Sampling")
         gr.HTML("""
         Runs <b><a href="https://huggingface.co/beyoru/Qwen3-0.9B-A0.6B" target="_blank">
-        beyoru/Qwen3-0.9B-A0.6B</a></b> via <b>Hugging Face Transformers</b>.<br><br>
-        <b>Supprot me at:</b>.<br><br>
+        beyoru/Qwen3-0.9B-A0.6B</a></b> with optional <b>Reasoning Sampling</b>.<br><br>
+        <b>Support me at:</b><br><br>
         <a href="https://www.buymeacoffee.com/ductransa0g" target="_blank">
-            <img src="https://cdn.buymeacoffee.com/buttons/v2/default-yellow.png" alt="Buy Me A Coffee" width="150px">
-          </a>
-        </p>
+            <img src="https://cdn.buymeacoffee.com/buttons/v2/default-yellow.png" 
+                 alt="Buy Me A Coffee" width="150px">
+        </a>
         """)
+        
         gr.Markdown("---")
-        gr.Markdown("## Generation Parameters")
+        gr.Markdown("## 🧠 Reasoning Settings")
+        
+        use_reasoning = gr.Checkbox(
+            label="Enable Reasoning Sampling",
+            value=False,
+            info="Sử dụng Metropolis-Hastings để cải thiện chất lượng"
+        )
+        
+        alpha = gr.Slider(
+            1.0, 5.0, value=2.0, step=0.1,
+            label="Alpha (Power)",
+            info="Độ 'sharp' của phân phối (càng cao càng tập trung vào token tốt nhất)"
+        )
+        
+        num_mcmc_steps = gr.Slider(
+            1, 20, value=5, step=1,
+            label="MCMC Steps per Token",
+            info="Số bước MCMC cho mỗi token (nhiều hơn = chất lượng cao hơn nhưng chậm hơn)"
+        )
+        
+        gr.Markdown("---")
+        gr.Markdown("## 📝 Standard Generation Parameters")
+        
         max_new_tokens = gr.Slider(32, 4096, value=1024, step=32, label="Max New Tokens")
         temperature = gr.Slider(0.1, 2.0, value=0.6, step=0.1, label="Temperature")
         repetition_penalty = gr.Slider(0.1, 2.0, value=1.0, step=0.1, label="Repetition Penalty")
         top_k = gr.Slider(0, 100, value=20, step=1, label="Top K (0 = off)")
         top_p = gr.Slider(0.0, 1.0, value=0.95, step=0.05, label="Top P")
-
+    
     gr.ChatInterface(
         fn=playground,
-        additional_inputs=[max_new_tokens, temperature, repetition_penalty, top_k, top_p],
+        additional_inputs=[
+            max_new_tokens, temperature, repetition_penalty, top_k, top_p,
+            use_reasoning, alpha, num_mcmc_steps
+        ],
         chatbot=gr.Chatbot(
-            label="Qwen3-0.9B-A0.6B",
+            label="Qwen3-0.9B-A0.6B with Reasoning",
             show_copy_button=True,
             allow_tags=["think"],
         ),
         examples=[
-            ["Hello who are you?"],
-            ["How to solve 2x+1=3."],
-            ["Example python code for async"]
+            ["Hello, who are you?"],
+            ["Solve the equation: 2x + 3 = 7"],
+            ["Write a Python function to calculate Fibonacci numbers"],
+            ["Explain quantum computing in simple terms"]
         ],
         cache_examples=False,
         show_api=False