Spaces:
Sleeping
Sleeping
File size: 2,270 Bytes
5b92639 c9dbbb0 5b92639 ee6cd87 5b92639 459508e 5b92639 c9dbbb0 5b92639 c9dbbb0 5b92639 5a88949 5b92639 5a88949 5b92639 c9dbbb0 5b92639 | 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 | import torch
import gradio as gr
from transformers import GPT2Tokenizer
from model import SmolLM2Config, SmolLM2ForCausalLM
tokenizer = GPT2Tokenizer.from_pretrained(
"HuggingFaceTB/cosmo2-tokenizer"
)
tokenizer.pad_token = tokenizer.eos_token
vocab_size = tokenizer.vocab_size
# Load the model
def load_model():
config = SmolLM2Config()
model = SmolLM2ForCausalLM(config) # Create base model instead of Lightning model
# Load just the model weights
state_dict = torch.load("model_weights.pth", map_location="cpu")['model_state_dict']
model.load_state_dict(state_dict)
model.eval()
return model
def generate_text(prompt, max_tokens, temperature=0.8, top_k=40):
"""Generate text based on the prompt"""
try:
# Encode the prompt
prompt_ids = tokenizer.encode(prompt, return_tensors="pt")
# Move to device if needed
device = next(model.parameters()).device
prompt_ids = prompt_ids.to(device)
# Generate text
with torch.no_grad():
generated_ids = model.generate( # Call generate directly on base model
prompt_ids,
max_new_tokens=max_tokens,
temperature=temperature,
top_k=top_k,
)
# Decode the generated text
generated_text = tokenizer.decode(generated_ids[0].tolist())
return generated_text
except Exception as e:
return f"An error occurred: {str(e)}"
# Load the model globally
model = load_model()
# Create the Gradio interface
demo = gr.Interface(
fn=generate_text,
inputs=[
gr.Textbox(
label="Enter your prompt", placeholder="Hello there!", lines=3
),
gr.Slider(
minimum=50,
maximum=500,
value=100,
step=10,
label="Maximum number of tokens",
),
],
outputs=gr.Textbox(label="Generated Text", lines=10),
title="Custom SmolLM2-135 Text Generator",
description="Enter a prompt and a custom SmolLM2 model will continue.",
examples=[
["The Nilakantha Series is a historically significant", 100],
["o find the second derivative of", 150],
],
)
if __name__ == "__main__":
demo.launch() |