Vesper27 commited on
Commit
ff29531
1 Parent(s): 935bae3

Create test.py

Browse files
Files changed (1) hide show
  1. test.py +24 -0
test.py ADDED
@@ -0,0 +1,24 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from transformers import AutoTokenizer
2
+ import transformers
3
+ import torch
4
+
5
+ model = "meta-llama/Llama-2-7b-chat-hf"
6
+
7
+ tokenizer = AutoTokenizer.from_pretrained(model)
8
+ pipeline = transformers.pipeline(
9
+ "text-generation",
10
+ model=model,
11
+ torch_dtype=torch.float16,
12
+ device_map="auto",
13
+ )
14
+
15
+ sequences = pipeline(
16
+ 'I liked "Breaking Bad" and "Band of Brothers". Do you have any recommendations of other shows I might like?\n',
17
+ do_sample=True,
18
+ top_k=10,
19
+ num_return_sequences=1,
20
+ eos_token_id=tokenizer.eos_token_id,
21
+ max_length=200,
22
+ )
23
+ for seq in sequences:
24
+ print(f"Result: {seq['generated_text']}")