File size: 7,088 Bytes
ad32177 06e1ad9 |
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 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 |
import gradio as gr
import os
from transformers import (
GPT2LMHeadModel, GPT2Tokenizer,
T5ForConditionalGeneration, T5Tokenizer,
AutoTokenizer, AutoModelForCausalLM
)
import torch
# Configuration for multiple models, can add more by extending MODEL_CONFIGS dict
MODEL_CONFIGS = {
"gpt2": {
"type": "causal",
"model_class": GPT2LMHeadModel,
"tokenizer_class": GPT2Tokenizer,
"description": "Original GPT-2, good for creative writing",
"size": "117M"
},
"distilgpt2": {
"type": "causal",
"model_class": AutoModelForCausalLM,
"tokenizer_class": AutoTokenizer,
"description": "Smaller, faster GPT-2",
"size": "82M"
},
"google/flan-t5-small": {
"type": "seq2seq",
"model_class": T5ForConditionalGeneration,
"tokenizer_class": T5Tokenizer,
"description": "Instruction-following T5 model",
"size": "80M"
},
"microsoft/DialoGPT-small": {
"type": "causal",
"model_class": AutoModelForCausalLM,
"tokenizer_class": AutoTokenizer,
"description": "Conversational AI model",
"size": "117M"
}
}
# Environment variables for optional authentication and private model access
HF_TOKEN = os.getenv("HF_TOKEN")
API_KEY = os.getenv("API_KEY")
ADMIN_PASSWORD = os.getenv("ADMIN_PASSWORD")
# Global state for caching loaded model and tokenizer
loaded_model_name = None
model = None
tokenizer = None
def load_model_and_tokenizer(model_name):
global loaded_model_name, model, tokenizer
if model_name == loaded_model_name and model is not None and tokenizer is not None:
return model, tokenizer
config = MODEL_CONFIGS[model_name]
if HF_TOKEN:
tokenizer = config["tokenizer_class"].from_pretrained(model_name, use_auth_token=HF_TOKEN)
model = config["model_class"].from_pretrained(model_name, use_auth_token=HF_TOKEN)
else:
tokenizer = config["tokenizer_class"].from_pretrained(model_name)
model = config["model_class"].from_pretrained(model_name)
# Set pad token for causal models if missing (important for generation padding)
if config["type"] == "causal" and tokenizer.pad_token is None:
tokenizer.pad_token = tokenizer.eos_token
loaded_model_name = model_name
return model, tokenizer
def authenticate_api_key(key):
if API_KEY and key != API_KEY:
return False
return True
def generate_text(prompt, model_name, max_length, temperature, top_p, top_k, api_key=""):
if API_KEY and not authenticate_api_key(api_key):
return "Error: Invalid API key"
try:
config = MODEL_CONFIGS[model_name]
model, tokenizer = load_model_and_tokenizer(model_name)
if config["type"] == "causal":
inputs = tokenizer.encode(prompt, return_tensors="pt", max_length=512, truncation=True)
with torch.no_grad():
outputs = model.generate(
inputs,
max_length=min(max_length + inputs.shape[1], 512),
temperature=temperature,
top_p=top_p,
top_k=top_k,
do_sample=True,
pad_token_id=tokenizer.pad_token_id,
num_return_sequences=1
)
generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
# Return generated continuation (remove original prompt)
return generated_text[len(prompt):].strip()
elif config["type"] == "seq2seq":
# Add task prefix for certain seq2seq models like flan-t5
task_prompt = f"Complete this text: {prompt}" if "flan-t5" in model_name.lower() else prompt
inputs = tokenizer(task_prompt, return_tensors="pt", max_length=512, truncation=True)
with torch.no_grad():
outputs = model.generate(
**inputs,
max_length=max_length,
temperature=temperature,
top_p=top_p,
top_k=top_k,
do_sample=True,
num_return_sequences=1
)
generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
return generated_text.strip()
except Exception as e:
return f"Error generating text: {str(e)}"
with gr.Blocks(title="Multi-Model Text Generation Server") as demo:
gr.Markdown("# Multi-Model Text Generation Server")
gr.Markdown("Choose a model from the dropdown, enter a text prompt, and generate text.")
with gr.Row():
with gr.Column():
model_selector = gr.Dropdown(
label="Model",
choices=list(MODEL_CONFIGS.keys()),
value="gpt2",
interactive=True
)
prompt_input = gr.Textbox(
label="Text Prompt",
placeholder="Enter the text prompt here...",
lines=4
)
max_length_slider = gr.Slider(
10, 200, 100, 10,
label="Max Generation Length"
)
temperature_slider = gr.Slider(
0.1, 2.0, 0.7, 0.1,
label="Temperature"
)
top_p_slider = gr.Slider(
0.1, 1.0, 0.9, 0.05,
label="Top-p (nucleus sampling)"
)
top_k_slider = gr.Slider(
1, 100, 50, 1,
label="Top-k sampling"
)
if API_KEY:
api_key_input = gr.Textbox(
label="API Key",
type="password",
placeholder="Enter API Key"
)
else:
api_key_input = gr.Textbox(value="", visible=False)
generate_btn = gr.Button("Generate Text", variant="primary")
with gr.Column():
output_textbox = gr.Textbox(
label="Generated Text",
lines=10,
placeholder="Generated text will appear here..."
)
generate_btn.click(
fn=generate_text,
inputs=[prompt_input, model_selector, max_length_slider, temperature_slider, top_p_slider, top_k_slider, api_key_input],
outputs=output_textbox
)
gr.Examples(
examples=[
["Once upon a time in a distant galaxy,"],
["The future of artificial intelligence is"],
["In the heart of the ancient forest,"],
["The detective walked into the room and noticed"],
],
inputs=prompt_input
)
auth_config = ("admin", ADMIN_PASSWORD) if ADMIN_PASSWORD else None
if __name__ == "__main__":
demo.launch(
auth=auth_config,
share=True, # Required for Spaces if localhost isn't accessible
server_name="0.0.0.0",
server_port=7860,
ssr_mode=False # Optional: disable server-side rendering to avoid Svelte i18n error
)
|