Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| from story_gpt.config import StoryGPTConfig | |
| from story_gpt.service import StoryGPTService | |
| config = StoryGPTConfig() | |
| service = StoryGPTService(config=config) | |
| def generate_story(title, genre, tone, idea, opening_line, max_new_tokens, temperature, top_k): | |
| return service.generate_story( | |
| title=title, | |
| genre=genre, | |
| tone=tone, | |
| idea=idea, | |
| opening_line=opening_line, | |
| max_new_tokens=int(max_new_tokens), | |
| temperature=float(temperature), | |
| top_k=int(top_k), | |
| ) | |
| def train_story_model(extra_story_text, steps): | |
| return service.train(extra_story_text=extra_story_text, steps=int(steps)) | |
| def reset_story_model(): | |
| return service.reset() | |
| with gr.Blocks( | |
| title="Story GPT Python", | |
| theme=gr.themes.Soft(primary_hue="amber", secondary_hue="orange"), | |
| ) as demo: | |
| gr.Markdown( | |
| """ | |
| # Story GPT Python | |
| A tiny story-writing GPT-style model written in Python from scratch. | |
| - Causal transformer decoder | |
| - Word-level tokenizer | |
| - Story-focused local training corpus | |
| - Structured local story composer for clean long-form output | |
| - No external pretrained LLM | |
| """ | |
| ) | |
| with gr.Tab("Write Story"): | |
| with gr.Row(): | |
| title_input = gr.Textbox(label="Title", value="The Intelligent Project") | |
| genre_input = gr.Dropdown( | |
| label="Genre", | |
| choices=[ | |
| "Fantasy", | |
| "Adventure", | |
| "Mystery", | |
| "Sci-Fi", | |
| "Friendship", | |
| "Folktale", | |
| "Educational", | |
| ], | |
| value="Educational", | |
| ) | |
| tone_input = gr.Dropdown( | |
| label="Tone", | |
| choices=["Warm", "Wonder", "Suspense", "Playful", "Calm", "Heroic", "Inspiring"], | |
| value="Inspiring", | |
| ) | |
| idea_input = gr.Textbox( | |
| label="Story Idea", | |
| value=( | |
| "A student builds an intelligent AI project step by step using Python, data analysis, " | |
| "machine learning, deep learning, and language models." | |
| ), | |
| lines=5, | |
| ) | |
| opening_line_input = gr.Textbox( | |
| label="Opening Line", | |
| value="Arman was a student who loved technology.", | |
| lines=2, | |
| ) | |
| with gr.Row(): | |
| max_tokens_input = gr.Slider(30, 220, value=110, step=5, label="Story Length") | |
| temperature_input = gr.Slider(0.2, 1.4, value=0.85, step=0.05, label="Temperature") | |
| top_k_input = gr.Slider(1, 24, value=10, step=1, label="Top-K") | |
| generate_button = gr.Button("Generate Story", variant="primary") | |
| output_text = gr.Textbox(label="Story Output", lines=14) | |
| output_status = gr.Textbox(label="Status", lines=4) | |
| with gr.Tab("Train"): | |
| extra_story_text_input = gr.Textbox( | |
| label="Extra Story Examples", | |
| placeholder="Add more short stories, story prompts, or endings to continue training the model.", | |
| lines=12, | |
| ) | |
| steps_input = gr.Slider(10, 500, value=140, step=10, label="Training Steps") | |
| train_button = gr.Button("Train Story Model", variant="primary") | |
| reset_button = gr.Button("Reset Model") | |
| train_status = gr.Textbox(label="Training Status", lines=6) | |
| generate_button.click( | |
| fn=generate_story, | |
| inputs=[ | |
| title_input, | |
| genre_input, | |
| tone_input, | |
| idea_input, | |
| opening_line_input, | |
| max_tokens_input, | |
| temperature_input, | |
| top_k_input, | |
| ], | |
| outputs=[output_text, output_status], | |
| ) | |
| train_button.click( | |
| fn=train_story_model, | |
| inputs=[extra_story_text_input, steps_input], | |
| outputs=[train_status], | |
| ) | |
| reset_button.click(fn=reset_story_model, outputs=[train_status]) | |
| if __name__ == "__main__": | |
| demo.launch() | |