Spaces:
Runtime error
Runtime error
File size: 3,578 Bytes
cddc4c2 fee88b4 cddc4c2 fee88b4 cddc4c2 fee88b4 cddc4c2 fee88b4 cddc4c2 fee88b4 cddc4c2 fee88b4 cddc4c2 fee88b4 cddc4c2 |
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 |
import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
# Load model and tokenizer
model_id = "jatingocodeo/SmolLM2"
def load_model():
try:
tokenizer = AutoTokenizer.from_pretrained(model_id)
# Ensure the tokenizer has the necessary special tokens
special_tokens = {
'pad_token': '[PAD]',
'eos_token': '</s>',
'bos_token': '<s>'
}
tokenizer.add_special_tokens(special_tokens)
model = AutoModelForCausalLM.from_pretrained(
model_id,
torch_dtype=torch.float16,
device_map="auto",
pad_token_id=tokenizer.pad_token_id
)
# Resize token embeddings to match new tokenizer
model.resize_token_embeddings(len(tokenizer))
return model, tokenizer
except Exception as e:
print(f"Error loading model: {str(e)}")
raise
def generate_text(prompt, max_length=100, temperature=0.7, top_k=50):
try:
# Load model and tokenizer (caching them for subsequent calls)
if not hasattr(generate_text, "model"):
generate_text.model, generate_text.tokenizer = load_model()
# Ensure the prompt is not empty
if not prompt.strip():
return "Please enter a prompt."
# Add BOS token if needed
if not prompt.startswith(generate_text.tokenizer.bos_token):
prompt = generate_text.tokenizer.bos_token + prompt
# Encode the prompt
input_ids = generate_text.tokenizer.encode(prompt, return_tensors="pt", truncation=True, max_length=2048)
input_ids = input_ids.to(generate_text.model.device)
# Generate text
with torch.no_grad():
output_ids = generate_text.model.generate(
input_ids,
max_length=min(max_length + len(input_ids[0]), 2048), # Respect model's max length
temperature=temperature,
top_k=top_k,
do_sample=True,
pad_token_id=generate_text.tokenizer.pad_token_id,
eos_token_id=generate_text.tokenizer.eos_token_id,
num_return_sequences=1
)
# Decode and return the generated text
generated_text = generate_text.tokenizer.decode(output_ids[0], skip_special_tokens=True)
return generated_text.strip()
except Exception as e:
print(f"Error during generation: {str(e)}")
return f"An error occurred: {str(e)}"
# Create Gradio interface
iface = gr.Interface(
fn=generate_text,
inputs=[
gr.Textbox(label="Prompt", placeholder="Enter your prompt here...", lines=2),
gr.Slider(minimum=10, maximum=200, value=100, step=1, label="Max Length"),
gr.Slider(minimum=0.1, maximum=1.0, value=0.7, step=0.1, label="Temperature"),
gr.Slider(minimum=1, maximum=100, value=50, step=1, label="Top K"),
],
outputs=gr.Textbox(label="Generated Text", lines=5),
title="SmolLM2 Text Generator",
description="""Generate text using the fine-tuned SmolLM2 model.
- Max Length: Controls the length of generated text
- Temperature: Controls randomness (higher = more creative)
- Top K: Controls diversity of word choices""",
examples=[
["Once upon a time", 100, 0.7, 50],
["The quick brown fox", 150, 0.8, 40],
["In a galaxy far far away", 200, 0.9, 30],
],
allow_flagging="never"
)
if __name__ == "__main__":
iface.launch() |