fittar commited on
Commit
98a1648
1 Parent(s): d3cb444

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +47 -0
app.py ADDED
@@ -0,0 +1,47 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+
3
+ from transformers import GPT2LMHeadModel, GPT2Tokenizer
4
+
5
+
6
+ def generate(text, model, tokenizer,device,do_sample,top_k=100, epsilon_cutoff=.00005, temperature=1):
7
+ #mark the text with special tokens
8
+ text=[tokenizer.eos_token + i + tokenizer.eos_token for i in text]
9
+ batch=tokenizer(text, padding=True, return_tensors="pt")
10
+
11
+ input_ids = batch["input_ids"].to(device)
12
+ attention_mask = batch["attention_mask"].to(device)
13
+
14
+ #how many new tokens to generate at max
15
+ max_prompt_length=50
16
+
17
+ generated_ids = model.generate(input_ids=input_ids,attention_mask=attention_mask, max_new_tokens=max_prompt_length, do_sample=do_sample,top_k=top_k, epsilon_cutoff=epsilon_cutoff, temperature=temperature)
18
+ #return only the generated prompts
19
+ pred_caps = tokenizer.batch_decode(generated_ids[:, -(generated_ids.shape[1] - input_ids.shape[1]):], skip_special_tokens=True)
20
+
21
+ return pred_caps
22
+
23
+ device = 'cuda' if torch.cuda.is_available() else 'cpu'
24
+ model = GPT2LMHeadModel.from_pretrained('fittar/ViPE-M-CTX7')
25
+ model.to(device)
26
+ tokenizer = GPT2Tokenizer.from_pretrained('gpt2-medium')
27
+ tokenizer.pad_token = tokenizer.eos_token
28
+
29
+
30
+ def generate(text):
31
+ result = gprompts=generate(text,model,tokenizer,do_sample=True,device=device)
32
+ return result
33
+
34
+ examples = [
35
+ ["Is string theory right?"],
36
+ ["She felt like a flower in December"],
37
+ ["2+2=4?"],
38
+ ]
39
+
40
+ demo = gr.Interface(
41
+ fn=generate,
42
+ inputs=gr.inputs.Textbox(lines=5, label="Input Text"),
43
+ outputs=gr.outputs.Textbox(label="Generated Text"),
44
+ examples=examples
45
+ )
46
+
47
+ demo.launch()