Spaces:
Sleeping
Sleeping
import gradio as gr | |
from transformers import AutoModelForCausalLM, AutoTokenizer | |
def load_model(model_name): | |
model = AutoModelForCausalLM.from_pretrained(model_name) | |
tokenizer = AutoTokenizer.from_pretrained(model_name) | |
return model, tokenizer | |
def converse(prompt, model1_name, model2_name, temperature1, temperature2): | |
model1, tokenizer1 = load_model(model1_name) | |
model2, tokenizer2 = load_model(model2_name) | |
inputs = tokenizer1(prompt, return_tensors="pt") | |
outputs = model1.generate(inputs.input_ids, max_new_tokens=50, temperature=temperature1, do_sample=True, pad_token_id=tokenizer1.eos_token_id) | |
response1 = tokenizer1.decode(outputs[0], skip_special_tokens=True) | |
inputs = tokenizer2(response1, return_tensors="pt") | |
outputs = model2.generate(inputs.input_ids, max_new_tokens=50, temperature=temperature2, do_sample=True, pad_token_id=tokenizer2.eos_token_id) | |
response2 = tokenizer2.decode(outputs[0], skip_special_tokens=True) | |
return response1, response2 | |
iface = gr.Interface( | |
fn=converse, | |
inputs=[ | |
gr.Textbox(label="Input Prompt"), | |
gr.Dropdown(["gpt2", "distilgpt2", "EleutherAI/gpt-neo-2.7B"], label="Model 1"), | |
gr.Dropdown(["gpt2", "distilgpt2", "EleutherAI/gpt-neo-2.7B"], label="Model 2"), | |
gr.Slider(0.1, 1.0, step=0.1, value=0.7, label="Temperature for Model 1"), | |
gr.Slider(0.1, 1.0, step=0.1, value=0.7, label="Temperature for Model 2") | |
], | |
outputs=["text", "text"], | |
title="Multi-Model Conversation", | |
description="Input a prompt to start a conversation between two models. Adjust temperatures for more diverse outputs." | |
) | |
if __name__ == "__main__": | |
iface.launch() | |