# Import libraries 
import gradio as gr
from transformers import pipeline
from huggingface_hub import InferenceClient
from huggingface_hub import InferenceApi
import numpy
import huggingface_hub 
import torch # assuming you're using PyTorch
import json
import pandas as pd
import sentencepiece
import accelerate
import safetensors
from transformers import T5Tokenizer, T5ForConditionalGeneration

# Set the model
tokenizer = T5Tokenizer.from_pretrained('Tianlin668/MentalT5')
model = T5ForConditionalGeneration.from_pretrained('Tianlin668/MentalT5')

# Generate text
def infer(prompt):
    input = f"<|startoftext|> {prompt.strip()}"
    input = tokenizer(input, return_tensors="pt").to(model.device)
    input_ids      = input["input_ids"]
    attention_mask = input["attention_mask"]

    output = model.generate(input_ids,
                            attention_mask=attention_mask,
                            max_new_tokens=100,
                            do_sample = True, top_k = 50, top_p = 0.85)
    output = tokenizer.decode(output[0], skip_special_tokens=True)
    return output
    
# Set interface
iface = gr.Interface(fn=infer, inputs=[gr.Textbox(label="Ask the Mental Doc Bot anything")], outputs="textbox")
if __name__ == "__main__":
    iface.launch()