james666600 commited on
Commit
cb95a71
·
verified ·
1 Parent(s): e3ac475

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +21 -9
app.py CHANGED
@@ -1,16 +1,28 @@
1
  import gradio as gr
2
- from transformers import pipeline
 
3
 
4
- # 直接加载 Hugging Face 模型
5
- classifier = pipeline("text-classification", model="bert-base-uncased")
 
 
 
 
 
 
 
6
 
7
- def predict(text):
8
- return classifier(text)
 
 
9
 
 
10
  demo = gr.Interface(
11
- fn=predict,
12
- inputs=gr.Textbox(label="输入文本"),
13
- outputs=gr.Label(label="分类结果"),
14
- title="BERT 文本分类演示",
15
  )
 
16
  demo.launch()
 
1
  import gradio as gr
2
+ from transformers import AutoModelForCausalLM, AutoTokenizer
3
+ import torch
4
 
5
+ # 加载 Qwen3-0.6B 模型和 tokenizer
6
+ model_name = "Qwen/Qwen3-0.6B"
7
+ tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
8
+ model = AutoModelForCausalLM.from_pretrained(
9
+ model_name,
10
+ torch_dtype=torch.float16, # 半精度减少显存占用
11
+ device_map="auto", # 自动选择 GPU
12
+ trust_remote_code=True # 信任远程代码(Qwen 需要)
13
+ )
14
 
15
+ def generate_text(prompt):
16
+ inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
17
+ outputs = model.generate(**inputs, max_new_tokens=100)
18
+ return tokenizer.decode(outputs[0], skip_special_tokens=True)
19
 
20
+ # 创建 Gradio 界面
21
  demo = gr.Interface(
22
+ fn=generate_text,
23
+ inputs=gr.Textbox(lines=3, placeholder="输入你的问题..."),
24
+ outputs=gr.Textbox(label="Qwen3-0.6B 的回答"),
25
+ title="Qwen3-0.6B 演示 (Free GPU)",
26
  )
27
+
28
  demo.launch()