from transformers import GPT2Tokenizer, GPT2LMHeadModel, TextDataset, DataCollatorForLanguageModeling

# Load the fine-tuned model and tokenizer
model = GPT2LMHeadModel.from_pretrained("kkhan/gpt2-medium-iba-faq")
tokenizer = GPT2Tokenizer.from_pretrained("kkhan/gpt2-medium-iba-faq")

def get_answer(model, tokenizer, question, max_length):
    input_ids = tokenizer.encode(question, return_tensors="pt")
    
    # Create the attention mask and pad token id
    attention_mask = torch.ones_like(input_ids)
    pad_token_id = tokenizer.eos_token_id

    output = model.generate(
        input_ids,
        max_length=max_length,
        num_return_sequences=1,
        attention_mask=attention_mask,
        pad_token_id=pad_token_id
    )

    return tokenizer.decode(output[0], skip_special_tokens=True)

 #Test the chatbot
#question = "Where does the IBA conduct aptitude tests and interviews?"  # Replace with your desired prompt
#answer = get_answer(model, tokenizer, question,50)
#print("Generated response:", answer)

import gradio as gr

def chat(chat_history, user_input):
  
  bot_response = get_answer(model, tokenizer,user_input,100)
  print(bot_response)
  response = ""
  for letter in ''.join(bot_response.response): #[bot_response[i:i+1] for i in range(0, len(bot_response), 1)]:
      response += letter + ""
      yield chat_history + [(user_input, bot_response)]

with gr.Blocks() as demo:
    gr.Markdown('# Q&A Bot with Fine-tuned Model')
    with gr.Tab("Knowledge Bot"):
#          inputbox = gr.Textbox("Input your text to build a Q&A Bot here.....")
          chatbot = gr.Chatbot()
          message = gr.Textbox ("what undergraduate programs does IBA offer?")
          message.submit(chat, [chatbot, message], chatbot)

demo.queue().launch(debug = True)