Bruno commited on
Commit
ed7e672
1 Parent(s): 4ec6c97

add app.py

Browse files
Files changed (1) hide show
  1. app.py +52 -0
app.py ADDED
@@ -0,0 +1,52 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from transformers import BlenderbotTokenizer, BlenderbotForConditionalGeneration
2
+ import torch
3
+ import gradio as gr
4
+
5
+ mname = "facebook/blenderbot-400M-distill"
6
+ model = BlenderbotForConditionalGeneration.from_pretrained(mname)
7
+ tokenizer = BlenderbotTokenizer.from_pretrained(mname)
8
+
9
+ def take_last_tokens(inputs, note_history, history):
10
+ """Filter the last 128 tokens"""
11
+ if inputs['input_ids'].shape[1] > 128:
12
+ inputs['input_ids'] = torch.tensor([inputs['input_ids'][0][-128:].tolist()])
13
+ inputs['attention_mask'] = torch.tensor([inputs['attention_mask'][0][-128:].tolist()])
14
+ note_history = ['</s> <s>'.join(note_history[0].split('</s> <s>')[2:])]
15
+ history = history[1:]
16
+ return inputs, note_history, history
17
+
18
+ def add_note_to_history(note, note_history):
19
+ """Add a note to the historical information"""
20
+ note_history.append(note)
21
+ note_history = '</s> <s>'.join(note_history)
22
+ return [note_history]
23
+
24
+ title = "Blenderbot Tokenizer"
25
+ description = """Blenderbot"""
26
+
27
+ def chat(message, history):
28
+ history = history or []
29
+ if history:
30
+ history_useful = ['</s> <s>'.join([str(a[0])+'</s> <s>'+str(a[1]) for a in history])]
31
+ else:
32
+ history_useful = []
33
+ history_useful = add_note_to_history(message, history_useful)
34
+ inputs = tokenizer(history_useful, return_tensors="pt")
35
+ inputs, history_useful, history = take_last_tokens(inputs, history_useful, history)
36
+ reply_ids = model.generate(**inputs)
37
+ response = tokenizer.batch_decode(reply_ids, skip_special_tokens=True)[0]
38
+ history_useful = add_note_to_history(response, history_useful)
39
+ list_history = history_useful[0].split('</s> <s>')
40
+ history.append((list_history[-2], list_history[-1]))
41
+ return history, history
42
+
43
+ gr.Interface(
44
+ fn=chat,
45
+ theme="huggingface",
46
+ css=".footer {display:none !important}",
47
+ inputs=["text", "state"],
48
+ outputs=["chatbot", "state"],
49
+ title=title,
50
+ description=description,
51
+ allow_flagging="never",
52
+ ).launch()