Spaces:
Running
Running
File size: 3,255 Bytes
617bd81 f586a0d 617bd81 f586a0d 617bd81 f586a0d 617bd81 f586a0d 617bd81 f586a0d |
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 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 |
import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
AVAILABLE_MODELS = {
"distilgpt2": "distilgpt2",
"bloomz-560m": "bigscience/bloomz-560m",
"gpt2-medium": "gpt2-medium",
"opt-350m": "facebook/opt-350m",
"pythia-160m": "EleutherAI/pythia-160m"
}
class TextGenerator:
def __init__(self):
self.model = None
self.tokenizer = None
def load_model(self, model_name: str) -> str:
try:
self.model = AutoModelForCausalLM.from_pretrained(AVAILABLE_MODELS[model_name])
self.tokenizer = AutoTokenizer.from_pretrained(AVAILABLE_MODELS[model_name])
return f"Successfully loaded {model_name}"
except Exception as e:
return f"Error loading model: {str(e)}"
def get_next_token_predictions(self, text: str, top_k: int = 10):
if not self.model or not self.tokenizer:
return [], []
inputs = self.tokenizer(text, return_tensors="pt")
with torch.no_grad():
outputs = self.model(**inputs)
logits = outputs.logits[0, -1, :]
probs = torch.nn.functional.softmax(logits, dim=-1)
top_k_probs, top_k_indices = torch.topk(probs, top_k)
top_k_tokens = [self.tokenizer.decode([idx.item()]) for idx in top_k_indices]
return top_k_tokens, top_k_probs.tolist()
generator = TextGenerator()
def format_predictions(tokens, probs):
if not tokens or not probs:
return "No predictions available"
formatted = "Predicted next tokens:\n\n"
for token, prob in zip(tokens, probs):
formatted += f"'{token}' : {prob:.4f}\n"
return formatted
def update_output(model_name, text, custom_token, selected_token):
output = text
if not generator.model or generator.model.name_or_path != AVAILABLE_MODELS[model_name]:
load_message = generator.load_model(model_name)
if "Error" in load_message:
return text, "", "", gr.update(choices=[]), load_message
if custom_token:
output += custom_token
elif selected_token:
output += selected_token.strip("'")
tokens, probs = generator.get_next_token_predictions(output)
predictions = format_predictions(tokens, probs)
token_choices = [f"'{token}'" for token in tokens]
return output, "", "", gr.update(choices=token_choices), predictions
demo = gr.Interface(
fn=update_output,
inputs=[
gr.Dropdown(choices=list(AVAILABLE_MODELS.keys()), value="distilgpt2", label="Select Model"),
gr.Textbox(lines=5, label="Generated Text", placeholder="Start typing or select a token..."),
gr.Textbox(label="Custom Token", placeholder="Type your own token..."),
gr.Dropdown(choices=[], label="Select from predicted tokens")
],
outputs=[
gr.Textbox(lines=5, label="Generated Text"),
gr.Textbox(label="Custom Token"),
gr.Textbox(label="Selected Token"),
gr.Dropdown(label="Predicted Tokens"),
gr.Textbox(lines=12, label="Predictions")
],
title="Interactive Text Generation",
description="Generate text by selecting predicted tokens or writing your own."
) |