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()