Improve story quality with structured composer and educational story mode

README.md

@@ -17,6 +17,7 @@ This is a tiny story-writing GPT-style language model project written in Python
 - Word-level tokenizer
 - Causal transformer decoder with self-attention
 - Story-focused local training corpus
+- Structured local story composer for cleaner paragraph output
 - Local CPU training loop
 - Checkpoint save and load
 - Gradio user interface

app.py

@@ -41,32 +41,44 @@ with gr.Blocks(
         - 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 Lantern in the Rain")
+            title_input = gr.Textbox(label="Title", value="The Intelligent Project")
             genre_input = gr.Dropdown(
                 label="Genre",
-                choices=["Fantasy", "Adventure", "Mystery", "Sci-Fi", "Friendship", "Folktale"],
-                value="Fantasy",
+                choices=[
+                    "Fantasy",
+                    "Adventure",
+                    "Mystery",
+                    "Sci-Fi",
+                    "Friendship",
+                    "Folktale",
+                    "Educational",
+                ],
+                value="Educational",
             )
             tone_input = gr.Dropdown(
                 label="Tone",
-                choices=["Warm", "Wonder", "Suspense", "Playful", "Calm", "Heroic"],
-                value="Wonder",
+                choices=["Warm", "Wonder", "Suspense", "Playful", "Calm", "Heroic", "Inspiring"],
+                value="Inspiring",
             )
 
         idea_input = gr.Textbox(
             label="Story Idea",
-            value="A child finds a glowing lantern that reveals hidden paths after a storm.",
+            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="When the rain stopped, the alley behind Mira's house began to shine.",
+            value="Arman was a student who loved technology.",
             lines=2,
         )
 

story_gpt/data.py

@@ -54,6 +54,14 @@ Tone: Calm
 Idea: A moon rabbit sews dreams into a blanket for a sleepless child.
 Opening: Above the sleeping city, one window was still lit.
 Story: Above the sleeping city, one window was still lit, and from the moon a rabbit noticed at once. The little rabbit gathered threads from cloud edges, silver dust from quiet stars, and a single feather from the night wind. It hopped down a beam of moonlight and sat on the child's windowsill, sewing without a sound. Into one square it stitched a forest with glowing fireflies. Into another it stitched a lake so still it could hold a whole sky. When the blanket touched the child, the room filled with the smell of rain and jasmine, and the tired eyes finally closed. By dawn the rabbit was back on the moon, but one silver thread remained on the pillow. The child kept it for years and slept as if the sky itself had tucked them in.
+
+Instruction: Write an inspiring educational story.
+Title: The Intelligent Project
+Genre: Educational
+Tone: Inspiring
+Idea: A student builds an intelligent AI project step by step using Python, data analysis, machine learning, deep learning, and language models.
+Opening: Arman was a student who loved technology.
+Story: Arman was a student who loved technology. One day, he decided to build an intelligent project using Artificial Intelligence and Machine Learning. At the beginning, he did not know much, but he was ready to learn. He started with Python and simple libraries such as NumPy and pandas so he could understand data clearly. After that, he used Matplotlib and Seaborn to visualize patterns and trends, which made his project easier to improve. Then he moved to Machine Learning with scikit-learn and built small predictive models step by step. Later, he explored Deep Learning with PyTorch and TensorFlow, and his project became powerful enough to work with images and text. In the final stage, he used Hugging Face tools so the system could answer questions and understand language. When he looked back at the finished project, Arman realized that the real success was not only the system he built, but the confidence he gained while building it.
 """.strip()
 
 

story_gpt/service.py

@@ -1,3 +1,4 @@
+import re
 import shutil
 
 import torch
@@ -27,6 +28,7 @@ class StoryGPTService:
         temperature: float,
         top_k: int,
     ):
+        clean_title = (title or "A New Story").strip()
         clean_prompt = build_story_prompt(
             title=title,
             genre=genre,
@@ -34,28 +36,49 @@ class StoryGPTService:
             idea=idea,
             opening_line=opening_line,
         )
-        self._ensure_ready()
-        encoded = self.tokenizer.encode(clean_prompt, add_bos=True)
-        idx = torch.tensor(encoded, dtype=torch.long).unsqueeze(0)
-        self.model.eval()
+        candidate_story = ""
+        used_mode = "structured"
 
-        with torch.inference_mode():
-            output = self.model.generate(
-                idx=idx,
+        if self.model is not None and self.tokenizer is not None:
+            candidate_story = self._generate_with_model(
+                prompt=clean_prompt,
                 max_new_tokens=max_new_tokens,
-                eos_id=self.tokenizer.eos_id,
                 temperature=temperature,
                 top_k=top_k,
             )
+            candidate_story = self._clean_generated_story(candidate_story)
+            if candidate_story and not self._looks_broken_story(candidate_story):
+                used_mode = "model"
+
+        if used_mode == "model":
+            story_text = self._format_story_output(clean_title, candidate_story)
+        else:
+            story_text = self._compose_story(
+                title=clean_title,
+                genre=genre,
+                tone=tone,
+                idea=idea,
+                opening_line=opening_line,
+            )
+
+        if self.model is None and self.config.checkpoint_path.exists():
+            self._ensure_ready()
+
+        if self.model is not None and self.tokenizer is not None:
+            model_info = (
+                f"Architecture=causal transformer, Vocab={self.tokenizer.vocab_size}, Layers={self.config.n_layers}."
+            )
+        else:
+            model_info = (
+                f"Architecture=causal transformer, Layers={self.config.n_layers}. "
+                f"Model training is available in the Train tab."
+            )
 
-        text = self.tokenizer.decode(output[0].tolist())
-        if "Story:" in text:
-            text = text.split("Story:", 1)[1].strip()
         status = (
-            f"Generated with Story GPT Python. "
-            f"Architecture=causal transformer, Vocab={self.tokenizer.vocab_size}, Layers={self.config.n_layers}."
+            f"Generated with Story GPT Python using {used_mode} generation. "
+            f"{model_info}"
         )
-        return text, status
+        return story_text, status
 
     def train(self, extra_story_text: str, steps: int):
         steps = max(1, steps)
@@ -96,6 +119,220 @@ class StoryGPTService:
             return
         self._load_or_initialize(extra_text="")
 
+    def _generate_with_model(self, prompt: str, max_new_tokens: int, temperature: float, top_k: int) -> str:
+        encoded = self.tokenizer.encode(prompt, add_bos=True)
+        idx = torch.tensor(encoded, dtype=torch.long).unsqueeze(0)
+        self.model.eval()
+
+        with torch.inference_mode():
+            output = self.model.generate(
+                idx=idx,
+                max_new_tokens=max_new_tokens,
+                eos_id=self.tokenizer.eos_id,
+                temperature=temperature,
+                top_k=top_k,
+            )
+
+        text = self.tokenizer.decode(output[0].tolist())
+        if "Story:" in text:
+            text = text.split("Story:", 1)[1].strip()
+        return text
+
+    def _clean_generated_story(self, text: str) -> str:
+        cleaned_lines = []
+        for raw_line in (text or "").splitlines():
+            line = raw_line.strip()
+            if not line:
+                continue
+            if re.match(r"^(Instruction|Title|Genre|Tone|Idea|Opening|Story)\s*:", line):
+                continue
+            cleaned_lines.append(line)
+
+        cleaned = " ".join(cleaned_lines)
+        cleaned = re.sub(r"\s+", " ", cleaned).strip()
+        cleaned = re.sub(r"([.!?])\1+", r"\1", cleaned)
+        return cleaned
+
+    def _looks_broken_story(self, text: str) -> bool:
+        if not text or len(text) < 220:
+            return True
+        if re.search(r"\b(Title|Genre|Tone|Idea|Opening|Story)\s*:", text):
+            return True
+
+        tokens = re.findall(r"[A-Za-z']+", text.lower())
+        if len(tokens) < 40:
+            return True
+
+        repeated_run = 1
+        current_run = 1
+        for prev, curr in zip(tokens, tokens[1:]):
+            if prev == curr:
+                current_run += 1
+                repeated_run = max(repeated_run, current_run)
+            else:
+                current_run = 1
+
+        unique_ratio = len(set(tokens)) / max(1, len(tokens))
+        sentence_count = len(re.findall(r"[.!?]", text))
+        return repeated_run >= 4 or unique_ratio < 0.45 or sentence_count < 5
+
+    def _format_story_output(self, title: str, story: str) -> str:
+        paragraphs = self._split_paragraphs(story)
+        if not paragraphs:
+            paragraphs = [story.strip()]
+        return f"Story: {title}\n\n" + "\n\n".join(paragraphs)
+
+    def _split_paragraphs(self, text: str):
+        sentences = re.split(r"(?<=[.!?])\s+", text.strip())
+        sentences = [sentence.strip() for sentence in sentences if sentence.strip()]
+        paragraphs = []
+        chunk = []
+
+        for sentence in sentences:
+            chunk.append(sentence)
+            if len(chunk) == 2:
+                paragraphs.append(" ".join(chunk))
+                chunk = []
+
+        if chunk:
+            paragraphs.append(" ".join(chunk))
+        return paragraphs
+
+    def _compose_story(self, title: str, genre: str, tone: str, idea: str, opening_line: str) -> str:
+        if self._is_project_story(title, idea, genre):
+            return self._compose_project_story(title, tone, idea, opening_line)
+        return self._compose_narrative_story(title, genre, tone, idea, opening_line)
+
+    def _is_project_story(self, title: str, idea: str, genre: str) -> bool:
+        text = f"{title} {idea} {genre}".lower()
+        keywords = [
+            "project",
+            "technology",
+            "ai",
+            "artificial intelligence",
+            "machine learning",
+            "deep learning",
+            "python",
+            "data",
+            "model",
+            "language",
+        ]
+        return genre == "Educational" or any(keyword in text for keyword in keywords)
+
+    def _compose_project_story(self, title: str, tone: str, idea: str, opening_line: str) -> str:
+        hero = self._choose_project_name(opening_line)
+        goal = self._project_goal(idea)
+        energy = {
+            "Inspiring": "with patience and belief",
+            "Heroic": "with determination and courage",
+            "Wonder": "with curiosity and excitement",
+            "Calm": "with steady focus",
+        }.get(tone, "with steady determination")
+
+        paragraphs = [
+            (
+                f"{opening_line.strip()} One day, {hero} decided to build {goal}. "
+                f"At the beginning, the path looked difficult, but {hero} moved forward {energy}."
+            ),
+            (
+                f"He started with Python and simple libraries such as NumPy and pandas so he could understand data clearly. "
+                f"Step by step, he learned how to collect, clean, and organize information, because he knew strong data is the base of every reliable system."
+            ),
+            (
+                f"After that, he used Matplotlib and Seaborn to visualize patterns and trends. "
+                f"The graphs showed him what the data was trying to say, and each chart made the next improvement easier to plan."
+            ),
+            (
+                f"Then he moved into Machine Learning with scikit-learn. "
+                f"He trained small models, checked their predictions, and improved them little by little until the results became more accurate and useful."
+            ),
+            (
+                f"Once the basics were strong, he explored Deep Learning with PyTorch and TensorFlow. "
+                f"Neural networks helped the project work with richer patterns in text and images, and the system started feeling more intelligent with every experiment."
+            ),
+            (
+                f"In the final stage, he connected modern language tools from Hugging Face so the project could answer questions, generate text, and understand language more naturally. "
+                f"When the work was finished, {hero} realized the greatest result was not only the intelligent system itself, but the confidence he had built while creating it."
+            ),
+        ]
+        return f"Story: {title}\n\n" + "\n\n".join(paragraphs)
+
+    def _choose_project_name(self, opening_line: str) -> str:
+        match = re.match(r"^([A-Z][a-z]+)\b", (opening_line or "").strip())
+        if match:
+            return match.group(1)
+        return "Arman"
+
+    def _project_goal(self, idea: str) -> str:
+        clean_idea = (idea or "").strip().rstrip(".")
+        if not clean_idea:
+            return "an intelligent project"
+        lowered = clean_idea.lower()
+        match = re.search(r"\bbuilds?\s+(.+)", lowered)
+        if match:
+            goal = clean_idea[match.start(1) :].strip().rstrip(".")
+            goal = re.sub(r"\s+using\s+.+$", "", goal, flags=re.IGNORECASE)
+            if goal:
+                if goal.lower().startswith(("a ", "an ", "the ")):
+                    return goal
+                return f"a {goal}"
+        if lowered.startswith("a ") or lowered.startswith("an ") or lowered.startswith("the "):
+            return clean_idea
+        return f"a project about {clean_idea}"
+
+    def _compose_narrative_story(
+        self,
+        title: str,
+        genre: str,
+        tone: str,
+        idea: str,
+        opening_line: str,
+    ) -> str:
+        scene = {
+            "Fantasy": "The world around that moment felt slightly magical, as if the ordinary street had made room for a secret.",
+            "Adventure": "It was the kind of beginning that quietly promises a journey.",
+            "Mystery": "Something about the air suggested that a hidden answer was waiting nearby.",
+            "Sci-Fi": "Even the smallest details seemed charged with future possibility.",
+            "Friendship": "It was the kind of moment that becomes important because someone chooses not to face it alone.",
+            "Folktale": "The elders would later say that such moments arrive only when the heart is ready.",
+            "Educational": "It was the kind of beginning that turns curiosity into steady progress.",
+        }.get(genre, "It was the kind of beginning that changes everything.")
+        mood = {
+            "Warm": "gently",
+            "Wonder": "with growing amazement",
+            "Suspense": "with careful attention",
+            "Playful": "with bright energy",
+            "Calm": "with quiet patience",
+            "Heroic": "with open courage",
+            "Inspiring": "with steady belief",
+        }.get(tone, "with steady purpose")
+        subject = (idea or "an unexpected challenge").strip().rstrip(".")
+        subject_text = subject[0].lower() + subject[1:] if subject else "an unexpected challenge"
+
+        paragraphs = [
+            (
+                f"{opening_line.strip()} {scene} The story truly began when the main character faced {subject_text} "
+                f"and chose to move forward {mood}."
+            ),
+            (
+                f"At first, the task felt larger than expected. "
+                f"Small clues, quiet observations, and a few brave decisions slowly revealed what needed to be done next."
+            ),
+            (
+                f"Along the way, each mistake turned into a lesson. "
+                f"The character learned to trust patience, notice details, and use both imagination and effort at the same time."
+            ),
+            (
+                f"When the hardest moment finally arrived, the answer came not from luck alone but from everything learned earlier. "
+                f"That was the point when uncertainty changed into action."
+            ),
+            (
+                f"By the end, the world around the character had shifted in a meaningful way. "
+                f"The problem was solved, but more importantly, the character had become wiser, stronger, and more ready for whatever came next."
+            ),
+        ]
+        return f"Story: {title}\n\n" + "\n\n".join(paragraphs)
+
     def _load_or_initialize(self, extra_text: str):
         checkpoint = self.config.checkpoint_path
         if checkpoint.exists():