import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer

vicuna_model = AutoModelForCausalLM.from_pretrained("lmsys/vicuna-7b-v1.3")
vicuna_tokenizer = AutoTokenizer.from_pretrained("lmsys/vicuna-7b-v1.3")

# llama_model = AutoModelForCausalLM.from_pretrained("luodian/llama-7b-hf")
# llama_tokenizer = AutoTokenizer.from_pretrained("luodian/llama-7b-hf")

# Define the function for generating responses
def generate_response(model, tokenizer, prompt):
    inputs = tokenizer(prompt, return_tensors="pt")
    outputs = model.generate(**inputs, max_length=500, pad_token_id=tokenizer.eos_token_id)
    response = tokenizer.decode(outputs[0])
    return response

# Define the Gradio interface
def chatbot_interface(prompt):
    vicuna_response = generate_response(vicuna_model, vicuna_tokenizer, prompt)
    # llama_response = generate_response(llama_model, llama_tokenizer, prompt)

    return {"Vicuna-7B": vicuna_response}

iface = gr.Interface(fn=chatbot_interface, 
                     inputs="text", 
                     outputs="text",
                     interpretation="default",
                     title="Chatbot with Three Models")

iface.launch()