File size: 2,051 Bytes
5782838 0f46be8 9d6743e 0f46be8 0343785 0f46be8 d7a34dd 0f46be8 0343785 d7a34dd 0f46be8 0343785 9d6743e 0f46be8 ca38751 0f46be8 e030ac0 0f46be8 0deb7d9 0f46be8 |
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 |
import gradio as gr
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer, TapasForQuestionAnswering, TapasTokenizer
# Load the models and tokenizers
tapas_model_name = "microsoft/tapex-large-finetuned-wtq"
dialogpt_model_name = "microsoft/DialoGPT-medium"
tapas_tokenizer = TapasTokenizer.from_pretrained(tapas_model_name)
tapas_model = BartForConditionalGeneration.from_pretrained(tapas_model_name)
dialogpt_tokenizer = AutoTokenizer.from_pretrained(dialogpt_model_name)
dialogpt_model = AutoModelForSeqCausalLM.from_pretrained(dialogpt_model_name)
def answer_table_question(table, question):
encoding = tapas_tokenizer(table=table, query=question, return_tensors="pt")
outputs = tapas_model.generate(**encoding)
response = tapas_tokenizer.batch_decode(outputs, skip_special_tokens=True)[0]
return response
def generate_dialog_response(prompt, conversation_history):
bot_input = dialogpt_tokenizer.encode(prompt + dialogpt_tokenizer.eos_token, return_tensors="pt")
chat_history_ids = dialogpt_model.generate(bot_input, max_length=1000, pad_token_id=dialogpt_tokenizer.eos_token_id)
response = dialogpt_tokenizer.decode(chat_history_ids[:, bot_input.shape[-1]:][0], skip_special_tokens=True)
return response
def chatbot_interface(user_input, table=gr.inputs.Textbox()):
global conversation_history
conversation_history.append(user_input)
# Check if user asks a question related to the table
if "table" in user_input:
question = user_input
answer = answer_table_question(table, question)
conversation_history.append(answer)
return "Bot (TAPAS): " + answer
else:
dialog_prompt = "User: " + " ".join(conversation_history) + "\nBot:"
response = generate_dialog_response(dialog_prompt, conversation_history)
conversation_history.append(response)
return "Bot (DialoGPT): " + response
conversation_history = []
iface = gr.Interface(fn=chatbot_interface, inputs=["text", "text"], outputs="text", live=True)
iface.launch() |