File size: 4,083 Bytes
6ec00a8
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
f82ff10
6ec00a8
 
 
 
 
 
f82ff10
6ec00a8
 
f82ff10
 
 
 
 
 
 
 
 
 
6ec00a8
 
 
f82ff10
 
6ec00a8
 
 
 
f82ff10
 
 
 
6ec00a8
 
 
 
f82ff10
6ec00a8
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
118
119
120
121
122
123
124
125
126
127
128
129
130
131
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()