app.py

app.py

@@ -0,0 +1,80 @@
+import os
+from typing import List, Tuple
+
+import gradio as gr
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+MODEL_ID = "Balab2021/qwen-workflow-planner-qwen2p5-lora"
+
+# Hugging Face Spaces automatically provides this if you set it in Secrets
+HF_TOKEN = os.getenv("HF_TOKEN")
+
+if not HF_TOKEN:
+    raise ValueError("HF_TOKEN environment variable is missing. Please add it in Space Settings → Secrets.")
+
+def build_messages(history: List[Tuple[str, str]], user_message: str):
+    messages = []
+    for user_text, assistant_text in history:
+        if user_text:
+            messages.append({"role": "user", "content": user_text})
+        if assistant_text:
+            messages.append({"role": "assistant", "content": assistant_text})
+    messages.append({"role": "user", "content": user_message})
+    return messages
+
+
+# Load model at startup
+print(f"Loading model: {MODEL_ID} ...")
+tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, token=HF_TOKEN)
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL_ID,
+    token=HF_TOKEN,
+    torch_dtype="auto",
+    device_map="auto",
+)
+
+
+def chat_fn(
+    message: str,
+    history: List[Tuple[str, str]],
+    temperature: float,
+    max_new_tokens: int,
+) -> str:
+    messages = build_messages(history, message)
+    
+    prompt = tokenizer.apply_chat_template(
+        messages,
+        tokenize=False,
+        add_generation_prompt=True,
+    )
+
+    inputs = tokenizer(prompt, return_tensors="pt")
+    inputs = {k: v.to(model.device) for k, v in inputs.items()}
+
+    with torch.no_grad():
+        output_ids = model.generate(
+            **inputs,
+            max_new_tokens=max_new_tokens,
+            temperature=temperature,
+            do_sample=temperature > 0,
+            pad_token_id=tokenizer.eos_token_id,
+        )
+
+    generated_ids = output_ids[0][inputs["input_ids"].shape[-1] :]
+    response = tokenizer.decode(generated_ids, skip_special_tokens=True).strip()
+    return response
+
+
+demo = gr.ChatInterface(
+    fn=chat_fn,
+    additional_inputs=[
+        gr.Slider(0.0, 1.5, value=0.2, step=0.05, label="Temperature"),
+        gr.Slider(32, 2048, value=512, step=32, label="Max New Tokens"),
+    ],
+    title="Qwen Workflow Planner Chat",
+    description=f"Model: {MODEL_ID}",
+)
+
+if __name__ == "__main__":
+    demo.launch()