Spaces:
Runtime error
Runtime error
File size: 2,407 Bytes
60c95e2 62e3b90 60c95e2 62e3b90 60c95e2 62e3b90 60c95e2 62e3b90 60c95e2 62e3b90 60c95e2 62e3b90 60c95e2 62e3b90 60c95e2 62e3b90 60c95e2 62e3b90 60c95e2 62e3b90 60c95e2 62e3b90 60c95e2 62e3b90 60c95e2 62e3b90 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 |
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
import gradio as gr
import concurrent.futures
# Set Device
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# Load Models
# Model 1: Bloom 560M
tokenizer1 = AutoTokenizer.from_pretrained('bigscience/bloom-560m')
model1 = AutoModelForCausalLM.from_pretrained('bigscience/bloom-560m', torch_dtype=torch.float16)
model1.to(device)
# Model 2: GPT-Neo 1.3B
tokenizer2 = AutoTokenizer.from_pretrained('EleutherAI/gpt-neo-1.3B')
model2 = AutoModelForCausalLM.from_pretrained('EleutherAI/gpt-neo-1.3B', torch_dtype=torch.float16)
model2.to(device)
# Define Functions with Improved Parameters
def generate_text_model1(prompt):
inputs = tokenizer1.encode(prompt, return_tensors='pt').to(device)
with torch.no_grad():
outputs = model1.generate(
inputs,
max_length=50,
num_return_sequences=1,
no_repeat_ngram_size=2,
do_sample=True,
top_k=50,
top_p=0.95,
temperature=0.8
)
text = tokenizer1.decode(outputs[0], skip_special_tokens=True)
return text
def generate_text_model2(prompt):
inputs = tokenizer2.encode(prompt, return_tensors='pt').to(device)
with torch.no_grad():
outputs = model2.generate(
inputs,
max_length=50,
num_return_sequences=1,
no_repeat_ngram_size=2,
do_sample=True,
top_k=50,
top_p=0.95,
temperature=0.8
)
text = tokenizer2.decode(outputs[0], skip_special_tokens=True)
return text
# Use ThreadPoolExecutor to Process in Parallel
def compare_models(prompt):
with concurrent.futures.ThreadPoolExecutor() as executor:
future1 = executor.submit(generate_text_model1, prompt)
future2 = executor.submit(generate_text_model2, prompt)
output1 = future1.result()
output2 = future2.result()
return output1, output2
# Create Gradio Interface
iface = gr.Interface(
fn=compare_models,
inputs=gr.Textbox(lines=2, placeholder='Enter a prompt here...'),
outputs=[gr.Textbox(label='Bloom 560M Output'), gr.Textbox(label='GPT-Neo 1.3B Output')],
title='Compare Text Generation Models',
description='Enter a prompt and see how two different models generate text.'
)
# Launch Interface
iface.launch()
|