Spaces:
Sleeping
Sleeping
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() |