import gradio as gr
from transformers import AutoTokenizer, TFBlenderbotForConditionalGeneration
import tensorflow as tf
import json
import os
from datetime import datetime

data = {"Interactions":[]}
with open("question_answer.json", "w") as file:
    json.dump(data, file, indent=4)

print("Loading the model......")
model_name = "WICKED4950/Esther_V1"
strategy = tf.distribute.MirroredStrategy()
tf.config.optimizer.set_jit(True)  # Enable XLA
tokenizer = AutoTokenizer.from_pretrained(model_name) 
with strategy.scope():
    model = TFBlenderbotForConditionalGeneration.from_pretrained(model_name)

def save_question(question,answer,path = "question_answer.json"):
    with open(path, "r") as file:
        data = json.load(file)
    data["Interactions"].append({"Question:":question,"Answer:":answer,"Time:":datetime.now().strftime("%Y-%m-%d %H:%M:%S")})
    with open(path, "w") as file:
        json.dump(data, file, indent=4)
        
print("Interface getting done....")
# Define the chatbot function
def predict(user_input):
    if user_input == "Print_data_hmm":
        with open("question_answer.json", "r") as file:
            print(json.load(file))
        return "Done"
    else:
        inputs = tokenizer(user_input, return_tensors="tf", padding=True, truncation=True,max_length=128)
    
        # Generate the response using the model
        response_id = model.generate(
            inputs['input_ids'],
            max_length=128,         # Set max length of response
            do_sample=True,         # Sampling for variability
            top_k=20,               # Consider top 50 tokens
            top_p=0.90,             # Nucleus sampling
            temperature=0.8         # Adjusts creativity of response
        )
    
        # Decode the response
        response = tokenizer.decode(response_id[0], skip_special_tokens=True)
        save_question(question = user_input,answer=response)
        return response

# Gradio interface
gr.Interface(
    fn=predict,
    inputs=gr.Textbox(label="Ask Esther anything!"),
    outputs=gr.Textbox(label="Esther's Response"),
    examples=[
        ["What should I do if I'm feeling down?"],
        ["How do I deal with stress?"],
        ["Tell me something positive!"]
    ],
     description="A chatbot trained to provide friendly and comforting responses. Type your question below and let Esther help!",
    title="Esther - Your Friendly Mental Health Chatbot",
    
).launch()