Spaces:
Sleeping
Sleeping
File size: 4,183 Bytes
1461fe5 | 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 | """
Gradio demo for the SLM-264M recipe generator fine-tune (ZeroGPU).
"""
import os
import spaces
import torch
import gradio as gr
from tokenizers import Tokenizer
from huggingface_hub import hf_hub_download
MODEL_REPO = os.environ.get("MODEL_REPO", "uday210/slm-recipe-264m") # <-- EDIT to your recipe repo
CKPT_FILENAME = os.environ.get("CKPT_FILENAME", "slm-recipe-final.pt") # <-- EDIT to real filename
TOKENIZER_FILENAME = os.environ.get("TOKENIZER_FILENAME", "tokenizer.json")
PROMPT_TEMPLATE = "Make a recipe using: {ingredients}\n"
BLOCK_SIZE = 2048
EOT_TOKEN = "<|endoftext|>"
from model import GPT # noqa: E402 (your import-safe model.py)
_token = os.environ.get("HF_TOKEN")
ckpt_path = hf_hub_download(MODEL_REPO, CKPT_FILENAME, token=_token)
tok_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), TOKENIZER_FILENAME)
ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=False)
model = GPT(ckpt["config"])
model.load_state_dict(ckpt["model"])
model.eval()
tok = Tokenizer.from_file(tok_path)
EOT_ID = tok.token_to_id(EOT_TOKEN)
_moved = False
@spaces.GPU(duration=120)
@torch.no_grad()
def generate(prompt: str, max_new_tokens: int = 400, temperature: float = 0.8, top_k: int = 50):
global _moved
if not _moved:
model.to("cuda")
_moved = True
ids = tok.encode(prompt).ids
x = torch.tensor(ids, dtype=torch.long, device="cuda")[None]
start = len(ids)
for _ in range(max_new_tokens):
logits, _ = model(x[:, -BLOCK_SIZE:])
logits = logits[:, -1, :]
if temperature and temperature > 0:
logits = logits / temperature
if top_k:
v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
logits[logits < v[:, [-1]]] = -float("inf")
probs = torch.softmax(logits, dim=-1)
nxt = torch.multinomial(probs, num_samples=1)
else:
nxt = logits.argmax(dim=-1, keepdim=True)
x = torch.cat([x, nxt], dim=1)
if EOT_ID is not None and nxt.item() == EOT_ID:
break
yield tok.decode(x[0, start:].tolist())
yield tok.decode(x[0, start:].tolist())
def make_recipe(ingredients, max_new_tokens, temperature):
ingredients = (ingredients or "").strip()
if not ingredients:
yield "Enter a few ingredients first (e.g. chicken, rice, lemon)."
return
prompt = PROMPT_TEMPLATE.format(ingredients=ingredients)
for partial in generate(prompt, int(max_new_tokens), float(temperature)):
yield partial
with gr.Blocks(title="SLM-264M Recipe Generator", theme=gr.themes.Soft()) as demo:
gr.Markdown(
"# 🍳 SLM-264M Recipe Generator\n"
"A 264M small language model, trained from scratch and fine-tuned to turn a list "
"of ingredients into a recipe.\n\n"
"*Note: format is reliable; it can occasionally repeat an ingredient or drift on "
"long recipes — it's a 264M model.*"
)
with gr.Row():
with gr.Column():
ingredients = gr.Textbox(
label="Ingredients", value="chicken, rice, lemon", lines=3,
placeholder="e.g. chicken, rice, lemon",
)
with gr.Accordion("Advanced", open=False):
max_tok = gr.Slider(100, 800, value=400, step=50, label="Max new tokens")
temp = gr.Slider(
0.0, 1.2, value=0.8, step=0.1,
label="Temperature (higher = more varied; ~0.8 is a good default)",
)
go = gr.Button("Generate recipe", variant="primary")
gr.Examples(
examples=[
["chicken, rice, lemon"],
["flour, sugar, butter, eggs, chocolate"],
["tomato, basil, garlic, pasta"],
],
inputs=[ingredients],
)
with gr.Column():
out = gr.Textbox(label="Recipe", lines=18, show_copy_button=True)
go.click(make_recipe, [ingredients, max_tok, temp], out)
ingredients.submit(make_recipe, [ingredients, max_tok, temp], out)
if __name__ == "__main__":
demo.launch() |