import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline

# Specify the directory containing the tokenizer's configuration file (config.json)
model_name = "pytorch_model-00001-of-00002.bin"

# Initialize the tokenizer
# tokenizer = AutoTokenizer.from_pretrained(model_name, local_files_only=True)
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
tokenizer.pad_token = tokenizer.eos_token
tokenizer.padding_side = "right"


# Initialize the GPT4All model
model = AutoModelForCausalLM.from_pretrained(model_name)

def generate_text(input_text):
    pipe = pipeline(task="text-generation", model=model, tokenizer=tokenizer, max_length=200)
    result = pipe(f"<s>[INST] {input_text} [/INST]")
    return result[0]['generated_text']

text_generation_interface = gr.Interface(
    fn=generate_text,
    inputs=[
        gr.inputs.Textbox(label="Input Text"),
    ],
    outputs=gr.outputs.Textbox(label="Generated Text"),
    title="GPT-4 Text Generation",
).launch()