tafxle commited on
Commit
59afef3
1 Parent(s): 1556b31

Create app.py

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