Tonic commited on
Commit
700caa0
1 Parent(s): 7339879

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +27 -38
app.py CHANGED
@@ -1,49 +1,38 @@
1
- from transformers import AutoModelForCausalLM, AutoTokenizer
2
- import gradio as gr
3
- import torch
4
-
5
- title = "👋🏻Welcome to Tonic's EZ Chat🚀"
6
- description = "You can use this Space to test out the current model (DialoGPT-medium) or duplicate this Space and use it for anyother model on 🤗HuggingFace."
7
- examples = [["How are you?"]]
8
-
9
- # Set the padding token to be used and initialize the model
10
- tokenizer = AutoTokenizer.from_pretrained("microsoft/DialoGPT-medium")
11
- tokenizer.padding_side = 'left'
12
-
13
  from transformers import AutoModelForCausalLM, AutoTokenizer
14
  import gradio as gr
15
  import torch
16
-
17
- title = "👋🏻Welcome to Tonic's EZ Chat🚀"
18
- description = "You can use this Space to test out the current model (DialoGPT-medium) or duplicate this Space and use it for any other model on 🤗HuggingFace. Join me on [Discord](https://discord.gg/fpEPNZGsbt) to build together."
19
- examples = [["How are you?"]]
20
-
21
  tokenizer = AutoTokenizer.from_pretrained("microsoft/DialoGPT-medium")
22
  tokenizer.padding_side = 'left'
23
- model = AutoModelForCausalLM.from_pretrained("microsoft/DialoGPT-medium")
24
 
25
- class ChatBot:
26
- def __init__(self):
27
- self.history = []
28
 
29
- def predict(self, input):
30
- new_user_input_ids = tokenizer.encode(input + tokenizer.eos_token, return_tensors="pt")
31
- bot_input_ids = torch.cat([torch.tensor(self.history), new_user_input_ids], dim=-1) if self.history else new_user_input_ids
32
- chat_history_ids = model.generate(bot_input_ids, max_length=4000, pad_token_id=tokenizer.eos_token_id)
33
- self.history.append(chat_history_ids[:, bot_input_ids.shape[-1]:].tolist())
34
- response = tokenizer.decode(chat_history_ids[:, bot_input_ids.shape[-1]:][0], skip_special_tokens=True)
 
35
  return response
36
 
37
- bot = ChatBot()
 
 
 
 
38
 
39
- iface = gr.Interface(
40
- fn=bot.predict,
41
- title=title,
42
- description=description,
43
- examples=examples,
44
- inputs="text",
45
- outputs="text",
46
- theme="ParityError/Anime",
47
- )
48
 
49
- iface.launch()
 
 
 
 
 
 
 
 
 
 
 
 
 
1
  from transformers import AutoModelForCausalLM, AutoTokenizer
2
  import gradio as gr
3
  import torch
4
+
 
 
 
 
5
  tokenizer = AutoTokenizer.from_pretrained("microsoft/DialoGPT-medium")
6
  tokenizer.padding_side = 'left'
7
+ model = AutoModelForCausalLM.from_pretrained("microsoft/DialoGPT-medium")
8
 
9
+ class ChatBot:
10
+ def __init__(self):
11
+ self.history = []
12
 
13
+ def predict(self, input):
14
+ new_user_input_ids = tokenizer.encode(input + tokenizer.eos_token, return_tensors="pt")
15
+ flat_history = [item for sublist in self.history for item in sublist]
16
+ bot_input_ids = torch.cat([torch.tensor(flat_history), new_user_input_ids], dim=-1) if self.history else new_user_input_ids
17
+ chat_history_ids = model.generate(bot_input_ids, max_length=2000, pad_token_id=tokenizer.eos_token_id)
18
+ self.history.append(chat_history_ids[:, bot_input_ids.shape[-1]:].tolist()[0])
19
+ response = tokenizer.decode(chat_history_ids[:, bot_input_ids.shape[-1]:][0], skip_special_tokens=True)
20
  return response
21
 
22
+ bot = ChatBot()
23
+
24
+ title = "👋🏻Welcome to Tonic's EZ Chat🚀"
25
+ description = "You can use this Space to test out the current model (DialoGPT-medium) or duplicate this Space and use it for any other model on 🤗HuggingFace."
26
+ examples = [["How are you?"]]
27
 
28
+ iface = gr.Interface(
29
+ fn=bot.predict,
30
+ title=title,
31
+ description=description,
32
+ examples=examples,
33
+ inputs="text",
34
+ outputs="text",
35
+ theme="ParityError/Anime"
36
+ )
37
 
38
+ iface.launch()