affandes commited on
Commit
acc073e
·
verified ·
1 Parent(s): f1dbd91

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +26 -32
app.py CHANGED
@@ -1,38 +1,32 @@
1
  from transformers import AutoModelForCausalLM, AutoTokenizer
2
- from fastapi import FastAPI
3
 
4
- app = FastAPI()
5
 
6
- @app.get("/")
7
- def greet_json():
8
- model_name = "Qwen/QwQ-32B-Preview"
 
 
 
9
 
10
- model = AutoModelForCausalLM.from_pretrained(
11
- model_name,
12
- torch_dtype="auto",
13
- device_map="auto"
14
- )
15
- tokenizer = AutoTokenizer.from_pretrained(model_name)
 
 
 
 
 
16
 
17
- prompt = "How many r in strawberry."
18
- messages = [
19
- {"role": "system", "content": "You are a helpful and harmless assistant. You are Qwen developed by Alibaba. You should think step-by-step."},
20
- {"role": "user", "content": prompt}
21
- ]
22
- text = tokenizer.apply_chat_template(
23
- messages,
24
- tokenize=False,
25
- add_generation_prompt=True
26
- )
27
- model_inputs = tokenizer([text], return_tensors="pt").to(model.device)
28
 
29
- generated_ids = model.generate(
30
- **model_inputs,
31
- max_new_tokens=512
32
- )
33
- generated_ids = [
34
- output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
35
- ]
36
-
37
- response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
38
- return response
 
1
  from transformers import AutoModelForCausalLM, AutoTokenizer
 
2
 
3
+ model_name = "Qwen/QwQ-32B-Preview"
4
 
5
+ model = AutoModelForCausalLM.from_pretrained(
6
+ model_name,
7
+ torch_dtype="auto",
8
+ device_map="auto"
9
+ )
10
+ tokenizer = AutoTokenizer.from_pretrained(model_name)
11
 
12
+ prompt = "How many r in strawberry."
13
+ messages = [
14
+ {"role": "system", "content": "You are a helpful and harmless assistant. You are Qwen developed by Alibaba. You should think step-by-step."},
15
+ {"role": "user", "content": prompt}
16
+ ]
17
+ text = tokenizer.apply_chat_template(
18
+ messages,
19
+ tokenize=False,
20
+ add_generation_prompt=True
21
+ )
22
+ model_inputs = tokenizer([text], return_tensors="pt").to(model.device)
23
 
24
+ generated_ids = model.generate(
25
+ **model_inputs,
26
+ max_new_tokens=512
27
+ )
28
+ generated_ids = [
29
+ output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
30
+ ]
 
 
 
 
31
 
32
+ response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]