Update app.py

app.py

@@ -1,74 +1,111 @@
-import gradio as gr
-from transformers import AutoModelForCausalLM, AutoTokenizer
+import os
 import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer
+from threading import Thread
+import gradio as gr
 
-# --- Load model ---
 MODEL_NAME = "beyoru/Qwen3-0.9B-A0.6B"
-tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+
+print("Loading model...")
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True)
 model = AutoModelForCausalLM.from_pretrained(
     MODEL_NAME,
-    torch_dtype=torch.float16,
+    torch_dtype=torch.bfloat16,
+    device_map="auto",
+    trust_remote_code=True
 )
+print("Model loaded.")
+
+
+def playground(
+    message,
+    history,
+    max_new_tokens,
+    temperature,
+    repetition_penalty,
+    top_k,
+    top_p
+):
+    if not isinstance(message, str) or not message.strip():
+        yield ""
+        return
+
+    # Build conversation
+    conversation = []
+    for user_msg, bot_msg in history:
+        conversation.append({"role": "user", "content": user_msg})
+        if bot_msg:
+            conversation.append({"role": "assistant", "content": bot_msg})
+    conversation.append({"role": "user", "content": message})
+
+    if hasattr(tokenizer, "apply_chat_template"):
+        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:"
 
-# --- Chat function ---
-def chat_fn(message, history, num_ctx, temperature, repeat_penalty, min_p, top_k, top_p, presence_penalty):
-    if not message.strip():
-        return ""
-    
-    # Tạo context chat từ lịch sử
-    conversation = ""
-    for turn in history:
-        role, content = turn["role"], turn["content"]
-        if role == "user":
-            conversation += f"User: {content}\n"
-        else:
-            conversation += f"Assistant: {content}\n"
-    conversation += f"User: {message}\nAssistant:"
-
-    inputs = tokenizer(conversation, return_tensors="pt", truncation=True, max_length=int(num_ctx)).to(model.device)
-
-    outputs = model.generate(
+    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+
+    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
+
+    generation_kwargs = dict(
         **inputs,
-        max_new_tokens=2048,
+        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),
-        top_k=int(top_k),
-        repetition_penalty=float(repeat_penalty),
-        do_sample=True,
-        eos_token_id=tokenizer.eos_token_id
+        repetition_penalty=float(repetition_penalty),
+        do_sample=True if temperature > 0 else False,
+        pad_token_id=tokenizer.eos_token_id
     )
 
-    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    # Cắt phần trước "Assistant:" để chỉ lấy câu trả lời
-    if "Assistant:" in response:
-        response = response.split("Assistant:")[-1].strip()
-    return response
+    # 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()
+
 
-# --- Giao diện Gradio ---
 with gr.Blocks(fill_height=True, fill_width=True) as app:
     with gr.Sidebar():
-        gr.Markdown("## Qwen3 Playground (Transformers Edition)")
-        gr.Markdown("Model: **beyoru/Qwen3-0.9B-A0.6B** — chạy trực tiếp bằng Transformers")
-
-        num_ctx = gr.Slider(512, 8192, 2048, 128, label="Context Length (num_ctx)")
-        temperature = gr.Slider(0.1, 2.0, 0.6, 0.1, label="Temperature")
-        repeat_penalty = gr.Slider(0.1, 2.0, 1.0, 0.1, label="Repeat Penalty")
-        min_p = gr.Slider(0.0, 1.0, 0.0, 0.01, label="Min P")
-        top_k = gr.Slider(0, 100, 20, 1, label="Top K")
-        top_p = gr.Slider(0.0, 1.0, 0.95, 0.05, label="Top P")
-        presence_penalty = gr.Slider(0.0, 2.0, 1.5, 0.1, label="Presence Penalty")
+        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> 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>
+        </p>
+        """)
+        gr.Markdown("---")
+        gr.Markdown("## Generation Parameters")
+        max_new_tokens = gr.Slider(32, 512, value=256, 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=chat_fn,
-        additional_inputs=[num_ctx, temperature, repeat_penalty, min_p, top_k, top_p, presence_penalty],
-        chatbot=gr.Chatbot(label="Transformers | Qwen3 (0.9B-A0.6B)", type="messages", show_copy_button=True),
+        fn=playground,
+        additional_inputs=[max_new_tokens, temperature, repetition_penalty, top_k, top_p],
+        chatbot=gr.Chatbot(
+            label="Qwen3-0.9B-A0.6B",
+            show_copy_button=True,
+            allow_tags=["think"],
+        ),
         examples=[
-            ["Introduce yourself."],
-            ["Explain quantum computers."],
-            ["Give a summary of World War II."]
+            ["Hello who are you?"],
+            ["How to solve 2x+1=3."],
+            ["Example python code for async"]
         ],
         cache_examples=False,
         show_api=False
     )
 
-app.launch(server_name="0.0.0.0", pwa=True)
+app.launch(server_name="0.0.0.0", pwa=True)