import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
import json

title = "AI ChatBot"
description = "A State-of-the-Art Large-scale Pretrained Response generation model (DialoGPT)"
examples = [["How are you?"]]

tokenizer = AutoTokenizer.from_pretrained("microsoft/DialoGPT-large")
model = AutoModelForCausalLM.from_pretrained("microsoft/DialoGPT-large")

def predict(input, history=[], file_path=None):
    if file_path:
        json_data = read_json_file(file_path)
        print(f"Contents of {file_path}:")
        print(json_data)
        print()
    new_user_input_ids = tokenizer.encode(input + tokenizer.eos_token, return_tensors="pt")
    bot_input_ids = torch.cat([torch.LongTensor(history), new_user_input_ids], dim=-1)
    history = model.generate(bot_input_ids, max_length=4000, pad_token_id=tokenizer.eos_token_id).tolist()
    response = tokenizer.decode(history[0]).split("\n")  # Splitting on new lines
    return response[0], history

def read_json_file(file_path):
    with open(file_path, 'r') as file:
        data = json.load(file)
    return data

def main():
    gr.Interface(
        fn=predict,
        title=title,
        description=description,
        examples=examples,
        inputs=[gr.inputs.Textbox(label="User Input"), gr.inputs.File(label="JSON File")],
        outputs=["text", "text"],
        theme="finlaymacklon/boxy_violet",
    ).launch()


if __name__ == "__main__":
    main()