kassemsabeh commited on
Commit
0997170
·
1 Parent(s): 4576fb4

Add application and requirements

Browse files
Files changed (2) hide show
  1. app.py +27 -0
  2. requirements.txt +2 -0
app.py ADDED
@@ -0,0 +1,27 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ import torch
3
+ from transformers import AutoTokenizer, T5ForConditionalGeneration
4
+
5
+ model_id = 'ksabeh/gavi'
6
+ max_input_length = 512
7
+ max_target_length = 10
8
+
9
+ model = T5ForConditionalGeneration.from_pretrained(model_id)
10
+ tokenizer = AutoTokenizer.from_pretrained(model_id)
11
+
12
+ def predict(title, category):
13
+ input = f"{title} <hl> {category} <hl>"
14
+ model_input = tokenizer(input, max_length=max_input_length, truncation=True,
15
+ padding="max_length")
16
+ model_input = {k:torch.unsqueeze(torch.tensor(v),dim=0) for k,v in model_input.items()}
17
+ predictions = model.generate(**model_input, num_beams=8, do_sample=True, max_length=10)
18
+ return tokenizer.batch_decode(predictions, skip_special_tokens=True)[0]
19
+
20
+ iface = gr.Interface(
21
+ predict,
22
+ inputs=["text", "text"],
23
+ outputs=['text'],
24
+ title="Attribute Generation",
25
+ )
26
+
27
+ iface.launch()
requirements.txt ADDED
@@ -0,0 +1,2 @@
 
 
 
1
+ transformers
2
+ torch