File size: 5,422 Bytes
04389c8
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
efc8ec6
04389c8
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
import torch
import gradio as gr
from gtts import gTTS
from transformers import pipeline
from transformers import GPT2LMHeadModel, GPT2Tokenizer
css = """
#input {background-color: #FFCCCB} 
"""
# Utility Functions
flatten = lambda l: [item for sublist in l for item in sublist]

def to_data(x):
    return x.data.numpy()

def to_var(x):
    if not torch.is_tensor(x):
        x = torch.Tensor(x)
    return x

def clear():
    return None,[],[]

def append(text, history,dialog_hx,personas):
    history.append([text,None])
    history , audio,dialog_hx= bot.respond(history,dialog_hx,personas)
    return history, audio, None,dialog_hx

class AI_Companion:
    """
    Class that Implements AI Companion.
    """

    def __init__(self, asr = "openai/whisper-tiny", chatbot = "af1tang/personaGPT"):
        """
        Create an Instance of the Companion.
        Parameters:
        asr: Huggingface ASR Model Card. Default: openai/whisper-tiny
        chatbot: Huggingface Conversational Model Card. Default: af1tang/personaGPT
        """

        self.device = torch.device('cpu')
        self.asr = pipeline("automatic-speech-recognition",model = asr,device= -1 if self.device == "cpu" else 0)
        self.model = GPT2LMHeadModel.from_pretrained(chatbot).to(self.device)
        self.tokenizer = GPT2Tokenizer.from_pretrained(chatbot)
        self.personas=[]
        self.sett={
            "do_sample":True,
            "top_k":10,
            "top_p":0.92,
            "max_length":1000,
        }

    def listen(self, audio, history):
        """
        Convert Speech to Text.

        Parameters:
        audio: Audio Filepath
        history: Chat History

        Returns:
        history : history with recognized text appended
        Audio : empty gradio component to clear gradio voice input
        """
        text = self.asr(audio)["text"]
        history.append([text,None])
        return history , None
    
    def add_fact(self,audio,personas,msg):
        '''
        Add fact to Persona.
        Takes in Audio, converts it into text and adds it to the facts list.

        Parameters:
        audio : audio of the spoken fact
        '''
        if audio is not None:
            text=self.asr(audio)
            personas.append(text['text']+self.tokenizer.eos_token)
        else:
            personas.append(msg+self.tokenizer.eos_token)
        return None,personas,None
    
    def respond(self, history,dialog_hx,personas,**kwargs):
        """
        Generates Response to User Input.

        Parameters:
        history: Chat History
        
        Returns:
        history: history with response appended
        audio: audio of the spoken response
        """

        person = self.tokenizer.encode(''.join(['<|p2|>'] + personas + ['<|sep|>'] + ['<|start|>']))
        user_inp= self.tokenizer.encode(history[-1][0]+self.tokenizer.eos_token)
        dialog_hx.append(user_inp)
        bot_input_ids = to_var([person + flatten(dialog_hx)]).long()
        with torch.no_grad():

            full_msg = self.model.generate(bot_input_ids,
                                        repetition_penalty=1.4,
                                        top_k = 10,
                                        top_p = 0.92,
                                        max_new_tokens = 256,
                                        num_beams=2,
                                        pad_token_id = self.tokenizer.eos_token_id)
        

        response = to_data(full_msg.detach()[0])[bot_input_ids.shape[-1]:]
        dialog_hx.append(response)
        history[-1][1] = self.tokenizer.decode(response, skip_special_tokens=True)
        self.speak(history[-1][1])
        return history, "out.mp3",dialog_hx
    
    def talk(self, audio, history,dialog_hx,personas,text):
        if audio is not None:
            history, _ = self.listen(audio, history)
        else:
            history.append([text,None])
        history, audio,dialog_hx = self.respond(history,dialog_hx,personas)
        return history, None, audio,dialog_hx,None

    def speak(self, text):
        """
        Speaks the given text using gTTS,
        Parameters:
        text: text to be spoken
        """
        tts = gTTS(text, lang='en')
        tts.save('out.mp3')

# Initialize AI Companion
bot = AI_Companion()
personas=[]
for i in ['I\'m a 19 year old girl','I study at IIT Indore','I am an easy-going and fun loving person','I love to swim','I am friendly, nice ,fun and kind','I am studious and get good grades']:
    response = i+ bot.tokenizer.eos_token
    personas.append(response)


# Create the Interface
with gr.Blocks() as demo:
    dialog_hx=gr.State([])
    personas=gr.State(personas)
    chatbot = gr.Chatbot([], elem_id = "chatbot").style(height = 300)
    audio = gr.Audio(source = "microphone", type = "filepath", label = "Input")
    msg = gr.Textbox()
    audio1 = gr.Audio(type = "filepath", label = "Output",elem_id="input")
    with gr.Row():
        b1 = gr.Button("Submit")
        b2 = gr.Button("Clear")
        b3=  gr.Button("Add Fact")
    b1.click(bot.talk, [audio, chatbot,dialog_hx,personas,msg], [chatbot, audio, audio1,dialog_hx,msg])
    msg.submit(append, [msg, chatbot,dialog_hx,personas], [chatbot, audio1, msg,dialog_hx])
    b2.click(clear, [] , [audio,chatbot,dialog_hx])
    b3.click(bot.add_fact, [audio,personas,msg], [audio,personas,msg])
demo.launch(share=True)