Emroi commited on
Commit
a2630cd
·
verified ·
1 Parent(s): c76546f

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +25 -0
app.py ADDED
@@ -0,0 +1,25 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from transformers import AutoTokenizer, AutoModelForCausalLM
2
+ import torch
3
+ model = AutoModelForCausalLM.from_pretrained("Vikhrmodels/Vikhr-7B-instruct_0.4",
4
+ device_map="auto",
5
+ attn_implementation="flash_attention_2",
6
+ torch_dtype=torch.bfloat16)
7
+
8
+ tokenizer = AutoTokenizer.from_pretrained("Vikhrmodels/Vikhr-7B-instruct_0.4")
9
+ from transformers import AutoTokenizer, pipeline
10
+ pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
11
+ prompts = [
12
+ "В чем разница между фруктом и овощем?",
13
+ "Годы жизни колмагорова?"]
14
+
15
+ def test_inference(prompt):
16
+ prompt = pipe.tokenizer.apply_chat_template([{"role": "user", "content": prompt}], tokenize=False, add_generation_prompt=True)
17
+ print(prompt)
18
+ outputs = pipe(prompt, max_new_tokens=512, do_sample=True, num_beams=1, temperature=0.25, top_k=50, top_p=0.98, eos_token_id=79097)
19
+ return outputs[0]['generated_text'][len(prompt):].strip()
20
+
21
+
22
+ for prompt in prompts:
23
+ print(f" prompt:\n{prompt}")
24
+ print(f" response:\n{test_inference(prompt)}")
25
+ print("-"*50)