File size: 918 Bytes
8cba1ca
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
import time
from transformers import pipeline, Conversation, AutoModelForCausalLM, AutoTokenizer
import gradio as gr

model_name = "microsoft/DialoGPT-large"

model = AutoModelForCausalLM.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name, padding_side='left')

chatbot = pipeline(task="conversational", model = model, tokenizer = tokenizer)

past_user_inputs = []
generated_responses = []

def chocolate_ai(message, history):
    conversation = Conversation(message, past_user_inputs = past_user_inputs, generated_responses = generated_responses)
    conversation = chatbot(conversation)

    past_user_inputs.append(message)
    generated_responses.append(conversation.generated_responses[-1])

    return conversation.generated_responses[-1]

gradio_interface = gr.ChatInterface(chocolate_ai, title="Chocolate AI", description="Type to start a conversation.")
gradio_interface.launch()