| """Streaming Gradio demo: python app.py --model checkpoints/release""" |
|
|
| import argparse |
| from pathlib import Path |
|
|
| import gradio as gr |
| import torch |
|
|
| from inference import FIELDS, load, stream |
|
|
|
|
| def build(model, tok): |
| def generate(summary, words, sentence, features, temperature, top_k, max_tokens): |
| values = (summary, words, sentence, features) |
| fields = dict(zip(FIELDS, ("" if value is None else value for value in values))) |
| try: |
| yield from stream(model, tok, fields, temperature, int(top_k), int(max_tokens)) |
| except (ValueError, TypeError, OverflowError) as e: |
| raise gr.Error(str(e)) from e |
|
|
| with gr.Blocks(title="Sprout · Story workshop") as demo: |
| gr.Markdown("# Sprout\nA small model learning to tell stories. Give it an idea and watch it write.") |
| gr.Markdown("Trained from scratch on TinyStories, then instruction tuned. " |
| "It can repeat itself, lose story details, and miss requested words or sentences.") |
| with gr.Row(): |
| with gr.Column(): |
| summary = gr.Textbox(label="Story idea", value="", lines=3, max_length=4096, |
| placeholder="A little fox plants a seed and learns to be patient.") |
| with gr.Accordion("Optional story constraints", open=False): |
| words = gr.Textbox(label="Words to include", value="", placeholder="fox, seed, patient", max_length=4096) |
| sentence = gr.Textbox(label="Sentence to include", value="", placeholder="Good things take time.", max_length=4096) |
| features = gr.Textbox(label="Features", value="", placeholder="Dialogue, MoralValue", max_length=4096) |
| with gr.Accordion("Generation settings", open=False): |
| temperature = gr.Slider(0, 2, value=0.8, step=0.05, label="Temperature", |
| info="0 gives the most likely next token; higher values add variety.") |
| top_k = gr.Slider(1, 200, value=50, step=1, label="Top-k") |
| max_tokens = gr.Slider(32, model.ctx, value=256, step=1, label="Maximum new tokens") |
| with gr.Row(): |
| write = gr.Button("Write a story", variant="primary") |
| stop = gr.Button("Stop") |
| with gr.Column(): |
| story = gr.Textbox(label="Your story", lines=18, interactive=False, |
| placeholder="Your story will appear here, one token at a time.") |
| gr.Examples(examples=[ |
| ["A little fox plants a seed and learns to be patient."], |
| ["Two friends find a red box beneath an old tree."], |
| ["A tiny dragon learns that blowing bubbles can help others."], |
| ], inputs=[summary]) |
| event = write.click(generate, [summary, words, sentence, features, temperature, top_k, max_tokens], |
| story, api_name="story", concurrency_limit=1) |
| stop.click(fn=None, cancels=[event], queue=False) |
| return demo.queue(max_size=8, default_concurrency_limit=1) |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser(description=__doc__) |
| local = Path(__file__).resolve().parent |
| bundle = local if (local / "config.json").is_file() else local / "checkpoints/release" |
| parser.add_argument("--model", default=str(bundle) if (bundle / "config.json").is_file() |
| else "GeorgeCodesSometimes/sprout-instruct") |
| parser.add_argument("--revision", help="Hub commit to load when --model is a model ID") |
| parser.add_argument("--host", help="Server bind address; defaults to Gradio's environment setting") |
| parser.add_argument("--port", type=int, default=7860) |
| args = parser.parse_args() |
| torch.set_num_threads(4) |
| model, tok, _ = load(args.model, revision=args.revision) |
| build(model, tok).launch(server_name=args.host, server_port=args.port, |
| theme=gr.themes.Default(primary_hue="green"), show_error=True) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
