| import gradio as gr |
| from transformers import pipeline |
| from transformers import Conversation |
|
|
| pipe = pipeline("conversational", model="PygmalionAI/pygmalion-6b") |
| conversation = Conversation() |
| conversation.add_message({"role": "assistant", "content": "How can I help you?"}) |
| memory = [] |
|
|
| def run_conversation(message, history): |
| |
| memory.append({"role": "user", "content": message}) |
| print(history) |
|
|
| if len(history) > 0: |
| for i in history: |
| conversation.add_message({"role":"user", "content":i[0]}) |
| conversation.add_message({"role": "assistant", "content": i[1]}) |
| for i in memory: |
| conversation.add_message(i) |
| print(conversation) |
|
|
| pipe(conversation) |
| return conversation.generated_responses[-1] |
|
|
| demo = gr.ChatInterface( |
| run_conversation, |
| chatbot = gr.Chatbot(height=500), |
| textbox = gr.Textbox(placeholder="Chat with me!", scale=7), |
| title = "Test", |
| description="Chat with me!", |
| examples=["hello"] |
| ) |
|
|
| if __name__ == "__main__": |
| demo.queue().launch() |
|
|
|
|
