jodiox commited on
Commit
2f62c52
·
verified ·
1 Parent(s): 3c7df73

update model to jodiox/olmo3-190m-zh-nano-sft

Browse files
Files changed (1) hide show
  1. app.py +38 -0
app.py ADDED
@@ -0,0 +1,38 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ from transformers import pipeline, AutoTokenizer
3
+
4
+ model_name = "jodiox/olmo3-190m-zh-nano-sft"
5
+
6
+ tokenizer = AutoTokenizer.from_pretrained(model_name)
7
+
8
+ # 收集结束标记 id
9
+ eos_token_ids = []
10
+ if tokenizer.eos_token_id is not None:
11
+ eos_token_ids.append(tokenizer.eos_token_id)
12
+ im_end_id = tokenizer.convert_tokens_to_ids("<|im_end|>")
13
+ if im_end_id is not None and im_end_id != tokenizer.unk_token_id:
14
+ eos_token_ids.append(im_end_id)
15
+
16
+ pipe = pipeline(
17
+ "text-generation",
18
+ model=model_name,
19
+ tokenizer=tokenizer,
20
+ model_kwargs={"attn_implementation": "sdpa"},
21
+ )
22
+
23
+ def predict(message):
24
+ messages = [{"role": "user", "content": message}]
25
+ output = pipe(
26
+ messages,
27
+ max_new_tokens=256,
28
+ do_sample=True,
29
+ temperature=0.7,
30
+ top_k=50,
31
+ top_p=0.9,
32
+ repetition_penalty=1.2,
33
+ return_full_text=False,
34
+ eos_token_id=eos_token_ids, # 关键:告诉 pipeline 遇到这些 token 就停
35
+ )
36
+ return output[0]["generated_text"]
37
+
38
+ gr.Interface(fn=predict, inputs="text", outputs="text").launch()