toymaker2 commited on
Commit
c025632
1 Parent(s): f3c5f2e

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +31 -0
app.py ADDED
@@ -0,0 +1,31 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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(
8
+ model,
9
+ use_auth_token=True,
10
+ )
11
+
12
+ pipeline = transformers.pipeline(
13
+ "text-generation",
14
+ model=model,
15
+ torch_dtype=torch.float16,
16
+ device_map="auto",
17
+ )
18
+
19
+ def gen(x, max_length=200):
20
+ sequences = pipeline(
21
+ x,
22
+ do_sample=True,
23
+ top_k=10,
24
+ num_return_sequences=1,
25
+ eos_token_id=tokenizer.eos_token_id,
26
+ max_length=max_length,
27
+ )
28
+
29
+ return sequences[0]["generated_text"].replace(x, "")
30
+
31
+ print(gen('I liked "Breaking Bad" and "Band of Brothers". Do you have any recommendations of other shows I might like?\n'))