LoufAn commited on
Commit
fe53d5f
·
1 Parent(s): 0747f6c

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +10 -13
app.py CHANGED
@@ -1,27 +1,24 @@
1
  import os
2
  import gradio as gr
 
3
  from transformers import AutoTokenizer, AutoModelForCausalLM, TextGenerationPipeline
4
 
5
- # 使用你在 Spaces 中添加的 Secret
6
  hf_token = os.environ.get("HUGGINGFACE_TOKEN")
7
 
8
- # 加载 tokenizer 和模型(带token)
9
- model_id = "google/gemma-3-4b-it"
10
- tokenizer = AutoTokenizer.from_pretrained(model_id, token=hf_token)
11
- model = AutoModelForCausalLM.from_pretrained(model_id, token=hf_token, device_map="auto")
12
-
13
- # 构造推理管道
14
- pipe = TextGenerationPipeline(model=model, tokenizer=tokenizer)
15
-
16
- # 推理函数
17
  def generate(prompt):
18
- output = pipe(prompt, max_new_tokens=100, do_sample=True, temperature=0.7)
 
 
 
19
  return output[0]["generated_text"]
20
 
21
- # Gradio UI
22
  gr.Interface(
23
  fn=generate,
24
  inputs=gr.Text(label="Enter your prompt"),
25
  outputs=gr.Textbox(label="Generated Text"),
26
- title="Gemma-3-27B Text Generation"
27
  ).launch()
 
1
  import os
2
  import gradio as gr
3
+ import spaces
4
  from transformers import AutoTokenizer, AutoModelForCausalLM, TextGenerationPipeline
5
 
6
+ model_id = "google/gemma-3-12b-it"
7
  hf_token = os.environ.get("HUGGINGFACE_TOKEN")
8
 
9
+ # 包含模型加载 + 推理
10
+ @spaces.GPU
 
 
 
 
 
 
 
11
  def generate(prompt):
12
+ tokenizer = AutoTokenizer.from_pretrained(model_id, token=hf_token)
13
+ model = AutoModelForCausalLM.from_pretrained(model_id, token=hf_token, device_map="auto")
14
+ pipeline = TextGenerationPipeline(model=model, tokenizer=tokenizer)
15
+ output = pipeline(prompt, max_new_tokens=100, do_sample=True, temperature=0.7)
16
  return output[0]["generated_text"]
17
 
18
+ # 构建界面
19
  gr.Interface(
20
  fn=generate,
21
  inputs=gr.Text(label="Enter your prompt"),
22
  outputs=gr.Textbox(label="Generated Text"),
23
+ title="Gemma-3-27B Inference (ZeroGPU)"
24
  ).launch()