from transformers import AutoModelForCausalLM, AutoTokenizer, AutoModelForSeq2SeqLM
tokenizer = AutoTokenizer.from_pretrained("microsoft/DialoGPT-large")
model = AutoModelForCausalLM.from_pretrained("microsoft/DialoGPT-large")
grammar_tokenizer = AutoTokenizer.from_pretrained("prithivida/grammar_error_correcter_v1")
grammar_model = AutoModelForSeq2SeqLM.from_pretrained("prithivida/grammar_error_correcter_v1")
import torch
import gradio as gr


def chat(message, history):
    history = history or []
    if message.startswith("How many"):
        response = random.randint(1, 10)
    elif message.startswith("How"):
        response = random.choice(["Great", "Good", "Okay", "Bad"])
    elif message.startswith("Where"):
        response = random.choice(["Here", "There", "Somewhere"])
    else:
        response = "I don't know"
    history.append((message, response))
    return history, feedback(message)


def feedback(text):
    tokenized_phrases = grammar_tokenizer([text], return_tensors='pt', padding=True)
    corrections = grammar_model.generate(**tokenized_phrases)
    corrections = grammar_tokenizer.batch_decode(corrections, skip_special_tokens=True)
    print("The corrections are: ", corrections)
    if corrections[0] == text:
        feedback = f'Looks good! Keep up the good work'
    else:
        feedback = f'\'{corrections[0]}\' might be a little better'
    return f'FEEDBACK:  {feedback}'

iface = gr.Interface(
    chat,
    ["text", "state"],
    ["chatbot", "text"],
    allow_screenshot=False,
    allow_flagging="never",
)
iface.launch()



new_user_input_ids = tokenizer.encode(text+tokenizer.eos_token, return_tensors='pt')
    # append the new user input tokens to the chat history
    bot_input_ids = torch.cat([chat_history_ids, new_user_input_ids], dim=-1) if chat_history_ids is not None else new_user_input_ids

    # generated a response while limiting the total chat history to 1000 tokens, 
    chat_history_ids = model.generate(bot_input_ids, max_length=5000, pad_token_id=tokenizer.eos_token_id)
    print("The text is ", [text])

    # pretty print last ouput tokens from bot
    output =  tokenizer.decode(chat_history_ids[:, bot_input_ids.shape[-1]:][0], skip_special_tokens=True)
    print("The outout is :", output)
    text_session.append(output)