Update app.py

app.py

@@ -1,12 +1,11 @@
 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 numpy as np
+import os
 
-MODEL_NAME = os.getenv('MODEL_ID', 'gpt2')
+MODEL_NAME = os.getenv('MODEL_ID')
 
 print("Loading model...")
 tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True)
@@ -18,125 +17,6 @@ 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,
@@ -145,15 +25,12 @@ def playground(
     temperature,
     repetition_penalty,
     top_k,
-    top_p,
-    use_reasoning,
-    alpha,
-    num_mcmc_steps
+    top_p
 ):
     if not isinstance(message, str) or not message.strip():
         yield ""
         return
-    
+
     # Build conversation
     conversation = []
     for user_msg, bot_msg in history:
@@ -161,122 +38,72 @@ 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])
-        prompt += "\nassistant:"
-    
-    # Setup streamer
+        prompt = "\n".join([f"{msg['role']}: {msg['content']}" for msg in conversation]) + "\nassistant:"
+
+    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+
     streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
-    
-    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
+
+    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)
     thread.start()
-    
+
     generated_text = ""
     for new_text in streamer:
         generated_text += new_text
         yield generated_text
-    
-    thread.join()
 
+    thread.join()
 
-# ===== GRADIO APP =====
 
 with gr.Blocks(fill_height=True, fill_width=True) as app:
     with gr.Sidebar():
-        gr.Markdown("## Playground with Reasoning Sampling")
+        gr.Markdown("## Playground by UltimaX Intelligence")
         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> with optional <b>Reasoning Sampling</b>.<br><br>
-        <b>Support me at:</b><br><br>
+        beyoru/Qwen3-0.9B-A0.6B</a></b> via <b>Hugging Face Transformers</b>.<br><br>
+        <b>Supprot 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>
+            <img src="https://cdn.buymeacoffee.com/buttons/v2/default-yellow.png" alt="Buy Me A Coffee" width="150px">
+          </a>
+        </p>
         """)
-        
         gr.Markdown("---")
-        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")
-        
+        gr.Markdown("## 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,
-            use_reasoning, alpha, num_mcmc_steps
-        ],
+        additional_inputs=[max_new_tokens, temperature, repetition_penalty, top_k, top_p],
         chatbot=gr.Chatbot(
-            label="Qwen3-0.9B-A0.6B with Reasoning",
+            label="Qwen3-0.9B-A0.6B",
             show_copy_button=True,
             allow_tags=["think"],
         ),
         examples=[
-            ["Hello, who are you?"],
-            ["Solve the equation: 2x + 3 = 7"],
-            ["Write a Python function to calculate Fibonacci numbers"],
-            ["Explain quantum computing in simple terms"]
+            ["Hello who are you?"],
+            ["How to solve 2x+1=3."],
+            ["Example python code for async"]
         ],
         cache_examples=False,
         show_api=False