codenlighten commited on
Commit
51d6cc4
·
verified ·
1 Parent(s): 1b190c4

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +27 -0
app.py ADDED
@@ -0,0 +1,27 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ from transformers import AutoModelForCausalLM, AutoTokenizer
3
+ import torch
4
+
5
+ # Load Model & Tokenizer
6
+ MODEL_NAME = "Qwen/Qwen2.5-3B" # Change to your HF model repo if fine-tuned
7
+
8
+ tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
9
+ model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, torch_dtype=torch.float16, device_map="auto")
10
+
11
+ def tinyzero_chat(prompt):
12
+ inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
13
+ outputs = model.generate(**inputs, max_length=512, temperature=0.7, top_p=0.9)
14
+ response = tokenizer.decode(outputs[0], skip_special_tokens=True)
15
+ return response
16
+
17
+ # Gradio Interface
18
+ demo = gr.Interface(
19
+ fn=tinyzero_chat,
20
+ inputs=gr.Textbox(placeholder="Ask TinyZero anything..."),
21
+ outputs=gr.Textbox(),
22
+ title="TinyZero Chatbot",
23
+ description="An interactive demo of TinyZero trained on reasoning tasks."
24
+ )
25
+
26
+ # Launch
27
+ demo.launch()