Waflon commited on
Commit
5be1431
·
verified ·
1 Parent(s): f2772e5

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +21 -0
app.py ADDED
@@ -0,0 +1,21 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from transformers import AutoTokenizer
2
+ import transformers
3
+ import torch
4
+ import accelerate
5
+
6
+ model = "meta-llama/Llama-2-7b-chat-hf"
7
+
8
+ tokenizer=AutoTokenizer.from_pretrained(model)
9
+ pipeline=transformers.pipeline(
10
+ "text-generation",
11
+ model=model,
12
+ tokenizer=tokenizer,
13
+ torch_dtype=torch.bfloat16,
14
+ trust_remote_code=True,
15
+ device_map="auto",
16
+ max_length=1000,
17
+ do_sample=True,
18
+ top_k=10,
19
+ num_return_sequences=1,
20
+ eos_token_id=tokenizer.eos_token_id
21
+ )