Update space with latest code and dependencies on Wed Jan 1 21:30:51 UTC 2025

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2023 Salvador Salazar
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

@@ -1,13 +1,24 @@
 ---
 title: Story Generator
-emoji: 🐨
-colorFrom: gray
-colorTo: red
+emoji: 📚
+colorFrom: blue
+python_version: 3.10.13
+colorTo: pink
 sdk: gradio
 sdk_version: 5.9.1
 app_file: app.py
 pinned: false
+preload_from_hub:
+# - meta-llama/Llama-3.2-1B-Instruct # Do not preload llama, as the token is not available at build time
+- google-bert/bert-base-uncased
+- facebook/bart-large-mnli
 license: mit
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+## Project Overview
+
+The Story Generator project leverages advanced natural language processing models to generate coherent and engaging stories. By utilizing models such as GPT-2, BERT, and BART, this project aims to provide users with a tool to create narratives based on given prompts. The application is built using Gradio for an interactive user interface, making it easy to input prompts and receive generated stories in real-time.
+
+The main purpose of this project is to explore the idea of beam search for selecting stories with high coherence, fluency, and genre alignment scores. This ensures that the generated stories are not only creative but also maintain a logical flow and adhere to the specified genre.
+
+Note that the final implementation is not strictly beam search and was modified to allow more diversity (creativity) inspired by the DVTS method in [this blog post](https://huggingface.co/spaces/HuggingFaceH4/blogpost-scaling-test-time-compute).

app.py

@@ -1,7 +1,136 @@
 import gradio as gr
+from typing import Literal
+from pydantic import BaseModel, Field, constr
+from story_beam_search.stories_generator import StoryGenerationSystem
+from typing import Tuple, List
 
-def greet(name):
-    return "Hello " + name + "!!"
+genre_choices = [
+    "children",
+    "mystery",
+    "adventure",
+    "sci-fi",
+    "fantasy",
+    "romance",
+    "comedy",
+    "drama",
+    "horror",
+]
 
-demo = gr.Interface(fn=greet, inputs="text", outputs="text")
-demo.launch()
+class InputModel(BaseModel):
+    prompt: str
+    genre: str
+    num_stories: int = Field(3, ge=2, le=7)
+    temperature: float = Field(2.5, ge=0.7, le=3.5)
+    max_length: int = Field(60, ge=30, le=200)
+
+
+def create_story_generation_interface() -> gr.Interface:
+    # Initialize the story generation system
+    system = StoryGenerationSystem()
+    system.initialize()
+
+    def generate_stories(
+        prompt: str, genre: str, num_stories: int, temperature: float, max_length: int
+    ) -> Tuple[str, List[str]]:
+        """
+        Generate and evaluate stories based on user input.
+        Returns a tuple of (detailed_scores, story_texts).
+        """
+
+        # Validate inputs.Gradio seems to validate chioces but not the range of the values
+        input_values = InputModel(
+            prompt=prompt, genre=genre, num_stories=num_stories, temperature=temperature, max_length=max_length
+        )
+
+        # Update beam search config with user parameters
+        system.beam_search.config.temperature = input_values.temperature
+        system.beam_search.config.max_length = input_values.max_length
+
+        # Generate and evaluate stories
+        ranked_stories = system.generate_and_evaluate(
+            input_values.prompt, input_values.genre, num_stories=input_values.num_stories
+        )
+
+        # Format detailed scores
+        detailed_scores = ""
+        story_texts = []
+
+        for i, (story, scores) in enumerate(ranked_stories, 1):
+            detailed_scores += f"Story {i}:\n"
+            detailed_scores += f"Total Score: {scores.total:.3f}\n"
+            detailed_scores += f"Coherence: {scores.coherence:.3f}\n"
+            detailed_scores += f"Fluency: {scores.fluency:.3f}\n"
+            detailed_scores += f"Genre Alignment: {scores.genre_alignment:.3f}\n"
+            detailed_scores += "-" * 50 + "\n"
+
+            story_texts.append(f"Story {i}:\n{story}\n")
+
+        return detailed_scores, "\n".join(story_texts)
+
+    # Define interface components
+    prompt_input = gr.Textbox(
+        label="Story Prompt",
+        placeholder="Enter the beginning of your story...",
+        lines=3,
+    )
+
+    genre_input = gr.Dropdown(
+        choices=genre_choices,
+        label="Genre",
+        value="fantasy",
+    )
+
+    num_stories_input = gr.Slider(
+        minimum=2, maximum=7, value=3, step=1, label="Number of Stories to Generate"
+    )
+
+    temperature_input = gr.Slider(
+        minimum=0.7, maximum=3.5, value=2.5, step=0.1, label="Temperature (Creativity)"
+    )
+
+    max_length_input = gr.Slider(
+        minimum=30, maximum=200, value=60, step=30, label="Maximum Length"
+    )
+
+    # Output components
+    scores_output = gr.Textbox(label="Detailed Scores", lines=10, interactive=False)
+
+    stories_output = gr.Textbox(label="Generated Stories", lines=15, interactive=False)
+
+    # Create the interface
+    interface = gr.Interface(
+        fn=generate_stories,
+        inputs=[
+            prompt_input,
+            genre_input,
+            num_stories_input,
+            temperature_input,
+            max_length_input,
+        ],
+        outputs=[scores_output, stories_output],
+        title="AI Story Generator",
+        description="""
+        Generate creative stories using AI! Enter a prompt and choose your preferences.
+        The system will generate multiple stories and evaluate them based on coherence,
+        fluency, and genre alignment.
+        """,
+        examples=[
+            ["Once upon a time in a magical forest,", "fantasy", 3, 1.8, 150],
+            [
+                "The detective knelt beside the bloodstained carpet, her gaze sharp as she traced the faint outline of a shoeprint.",
+                "mystery",
+                3,
+                2.7,
+                200,
+            ],
+        ],
+        theme=gr.themes.Soft(),
+    )
+
+    return interface
+
+
+if __name__ == "__main__":
+    # Create and launch the interface
+    interface = create_story_generation_interface()
+    interface.launch(show_error=True)

requirements.txt

@@ -0,0 +1,79 @@
+aiofiles==23.2.1 ; python_full_version == "3.10.13"
+annotated-types==0.7.0 ; python_full_version == "3.10.13"
+anyio==4.7.0 ; python_full_version == "3.10.13"
+certifi==2024.12.14 ; python_full_version == "3.10.13"
+charset-normalizer==3.4.1 ; python_full_version == "3.10.13"
+click==8.1.8 ; sys_platform != "emscripten" and python_full_version == "3.10.13"
+colorama==0.4.6 ; python_full_version == "3.10.13" and platform_system == "Windows"
+exceptiongroup==1.2.2 ; python_full_version == "3.10.13"
+fastapi==0.115.6 ; python_full_version == "3.10.13"
+ffmpy==0.5.0 ; python_full_version == "3.10.13"
+filelock==3.16.1 ; python_full_version == "3.10.13"
+fsspec==2024.12.0 ; python_full_version == "3.10.13"
+gradio-client==1.5.2 ; python_full_version == "3.10.13"
+gradio==5.9.1 ; python_full_version == "3.10.13"
+h11==0.14.0 ; python_full_version == "3.10.13"
+httpcore==1.0.7 ; python_full_version == "3.10.13"
+httpx==0.28.1 ; python_full_version == "3.10.13"
+huggingface-hub==0.27.0 ; python_full_version == "3.10.13"
+idna==3.10 ; python_full_version == "3.10.13"
+jinja2==3.1.5 ; python_full_version == "3.10.13"
+joblib==1.4.2 ; python_full_version == "3.10.13"
+markdown-it-py==3.0.0 ; sys_platform != "emscripten" and python_full_version == "3.10.13"
+markupsafe==2.1.5 ; python_full_version == "3.10.13"
+mdurl==0.1.2 ; sys_platform != "emscripten" and python_full_version == "3.10.13"
+mpmath==1.3.0 ; python_full_version == "3.10.13"
+networkx==3.4.2 ; python_full_version == "3.10.13"
+numpy==2.2.1 ; python_full_version == "3.10.13"
+nvidia-cublas-cu12==12.1.3.1 ; platform_system == "Linux" and platform_machine == "x86_64" and python_full_version == "3.10.13"
+nvidia-cuda-cupti-cu12==12.1.105 ; platform_system == "Linux" and platform_machine == "x86_64" and python_full_version == "3.10.13"
+nvidia-cuda-nvrtc-cu12==12.1.105 ; platform_system == "Linux" and platform_machine == "x86_64" and python_full_version == "3.10.13"
+nvidia-cuda-runtime-cu12==12.1.105 ; platform_system == "Linux" and platform_machine == "x86_64" and python_full_version == "3.10.13"
+nvidia-cudnn-cu12==9.1.0.70 ; platform_system == "Linux" and platform_machine == "x86_64" and python_full_version == "3.10.13"
+nvidia-cufft-cu12==11.0.2.54 ; platform_system == "Linux" and platform_machine == "x86_64" and python_full_version == "3.10.13"
+nvidia-curand-cu12==10.3.2.106 ; platform_system == "Linux" and platform_machine == "x86_64" and python_full_version == "3.10.13"
+nvidia-cusolver-cu12==11.4.5.107 ; platform_system == "Linux" and platform_machine == "x86_64" and python_full_version == "3.10.13"
+nvidia-cusparse-cu12==12.1.0.106 ; platform_system == "Linux" and platform_machine == "x86_64" and python_full_version == "3.10.13"
+nvidia-nccl-cu12==2.20.5 ; platform_system == "Linux" and platform_machine == "x86_64" and python_full_version == "3.10.13"
+nvidia-nvjitlink-cu12==12.6.85 ; platform_system == "Linux" and platform_machine == "x86_64" and python_full_version == "3.10.13"
+nvidia-nvtx-cu12==12.1.105 ; platform_system == "Linux" and platform_machine == "x86_64" and python_full_version == "3.10.13"
+orjson==3.10.13 ; python_full_version == "3.10.13"
+packaging==24.2 ; python_full_version == "3.10.13"
+pandas==2.2.3 ; python_full_version == "3.10.13"
+pillow==11.0.0 ; python_full_version == "3.10.13"
+protobuf==5.29.2 ; python_full_version == "3.10.13"
+pydantic-core==2.27.2 ; python_full_version == "3.10.13"
+pydantic==2.10.4 ; python_full_version == "3.10.13"
+pydub==0.25.1 ; python_full_version == "3.10.13"
+pygments==2.18.0 ; sys_platform != "emscripten" and python_full_version == "3.10.13"
+python-dateutil==2.9.0.post0 ; python_full_version == "3.10.13"
+python-multipart==0.0.20 ; python_full_version == "3.10.13"
+pytz==2024.2 ; python_full_version == "3.10.13"
+pyyaml==6.0.2 ; python_full_version == "3.10.13"
+regex==2024.11.6 ; python_full_version == "3.10.13"
+requests==2.32.3 ; python_full_version == "3.10.13"
+rich==13.9.4 ; sys_platform != "emscripten" and python_full_version == "3.10.13"
+ruff==0.8.4 ; sys_platform != "emscripten" and python_full_version == "3.10.13"
+safehttpx==0.1.6 ; python_full_version == "3.10.13"
+safetensors==0.4.5 ; python_full_version == "3.10.13"
+scikit-learn==1.6.0 ; python_full_version == "3.10.13"
+scipy==1.14.1 ; python_full_version == "3.10.13"
+semantic-version==2.10.0 ; python_full_version == "3.10.13"
+shellingham==1.5.4 ; sys_platform != "emscripten" and python_full_version == "3.10.13"
+six==1.17.0 ; python_full_version == "3.10.13"
+sniffio==1.3.1 ; python_full_version == "3.10.13"
+starlette==0.41.3 ; python_full_version == "3.10.13"
+sympy==1.13.3 ; python_full_version == "3.10.13"
+threadpoolctl==3.5.0 ; python_full_version == "3.10.13"
+tokenizers==0.21.0 ; python_full_version == "3.10.13"
+tomlkit==0.13.2 ; python_full_version == "3.10.13"
+torch==2.4.0 ; python_full_version == "3.10.13"
+tqdm==4.67.1 ; python_full_version == "3.10.13"
+transformers==4.47.1 ; python_full_version == "3.10.13"
+triton==3.0.0 ; platform_system == "Linux" and platform_machine == "x86_64" and python_full_version == "3.10.13"
+typer==0.15.1 ; sys_platform != "emscripten" and python_full_version == "3.10.13"
+typing-extensions==4.12.2 ; python_full_version == "3.10.13"
+tzdata==2024.2 ; python_full_version == "3.10.13"
+urllib3==2.3.0 ; python_full_version == "3.10.13"
+uvicorn==0.34.0 ; sys_platform != "emscripten" and python_full_version == "3.10.13"
+websockets==14.1 ; python_full_version == "3.10.13"

story_beam_search/beam_search.py

@@ -0,0 +1,103 @@
+from pydantic.dataclasses import dataclass
+from typing import Optional
+import torch
+from transformers import PreTrainedModel, PreTrainedTokenizer
+
+from story_beam_search.scoring import StoryScorer
+
+
+@dataclass
+class BeamSearchConfig:
+    num_beams: int = 3
+    num_return_sequences: int = 3
+    max_length: int = 100
+    no_repeat_ngram_size: int = 2
+    temperature: float = 0.8
+    top_k: int = 8
+    top_p: float = 0.95
+    num_iterations: int = 3
+    continuation_length: int = 10
+
+
+class BeamSearchGenerator:
+    def __init__(
+        self,
+        model: PreTrainedModel,
+        tokenizer: PreTrainedTokenizer,
+        device: torch.device,
+        config: Optional[BeamSearchConfig] = None,
+    ):
+        self.model = model
+        self.tokenizer = tokenizer
+        self.device = device
+        self.config = config or BeamSearchConfig()
+
+    def generate_iterations(
+        self, prompt: str, genre: str, evaluator: StoryScorer
+    ) -> list[str]:
+        """
+        Generate story continuations using multiple iterations of beam search.
+        """
+
+        # Adding some instructions to the prompt. These are removed in the end
+        instructions = (
+            f"Continue the following story in the {genre} genre, "
+            "ensuring coherence with the tone, characters, and narrative established so far:\n"
+        )
+        instructions_len = len(instructions)
+
+        stories = self._generate_single_iteration(instructions + prompt)
+        ranked_stories = evaluator.evaluate_multiple(
+            [story[instructions_len:] for story in stories]
+        )
+
+        stories = [story for story, _ in ranked_stories[:self.config.num_beams]]
+
+        if stories:
+            for _ in range(self.config.num_iterations):
+                all_stories = []
+                for story in stories:
+                    continuations = self._generate_single_iteration(
+                        instructions + story
+                    )
+                    all_stories.extend(continuations)
+                ranked_stories = evaluator.evaluate_multiple(
+                    [story[instructions_len:] for story in all_stories]
+                )
+                stories = [story for story, _ in ranked_stories[:self.config.num_beams]]
+
+        return stories
+
+    def _generate_single_iteration(self, prompt: str) -> list[str]:
+        """
+        Generate multiple continuations for a single iteration using beam search.
+        """
+        inputs = self.tokenizer(prompt, return_tensors="pt").to(self.device)
+        input_ids = inputs["input_ids"]
+        attention_mask = inputs["attention_mask"]
+
+        self.config.continuation_length = (
+            len(input_ids[0]) + self.config.max_length // self.config.num_iterations
+        )
+
+        with torch.no_grad():
+            outputs = self.model.generate(
+                input_ids=input_ids,
+                attention_mask=attention_mask,
+                max_length=self.config.continuation_length,
+                num_beams=self.config.num_beams,
+                num_return_sequences=self.config.num_return_sequences,
+                early_stopping=True,
+                no_repeat_ngram_size=self.config.no_repeat_ngram_size,
+                temperature=self.config.temperature,
+                top_k=self.config.top_k,
+                top_p=self.config.top_p,
+                do_sample=True,
+            )
+
+        stories = []
+        for output in outputs:
+            text = self.tokenizer.decode(output, skip_special_tokens=True)
+            stories.append(text)
+
+        return stories

story_beam_search/scoring.py

@@ -0,0 +1,180 @@
+import numpy as np
+from dataclasses import dataclass
+from typing import Protocol
+import torch
+from transformers import PreTrainedModel, PreTrainedTokenizer, Pipeline
+from sklearn.metrics.pairwise import cosine_similarity
+
+
+class StoryScorer(Protocol):
+    """Protocol defining the interface for story scoring components."""
+
+    def score(self, story: str) -> float:
+        """Return a score between 0 and 1."""
+        ...
+
+
+@dataclass
+class CombinedScore:
+    coherence: float = 0.0
+    fluency: float = 0.0
+    genre_alignment: float = 0.0
+    total: float = 0.0
+
+
+class CoherenceScorer(StoryScorer):
+    def __init__(
+        self,
+        model: PreTrainedModel,
+        tokenizer: PreTrainedTokenizer,
+        device: torch.device,
+        max_pairs: int = 3,
+    ):
+        self.model = model
+        self.tokenizer = tokenizer
+        self.device = device
+        self.max_pairs = max_pairs
+
+    def score(self, story: str) -> float:
+        """Calculate coherence score based on sentences cosine similarity."""
+
+        sentences = [s.strip() for s in story.split(".") if s.strip()]
+
+        embeddings = []
+
+        # Generate embeddings for each sentence
+        for sentence in sentences:
+            inputs = self.tokenizer(sentence, return_tensors="pt").to(self.device)
+            with torch.no_grad():
+                emb = self.model.bert(**inputs).last_hidden_state[:, 0, :]
+                embeddings.append(emb.cpu().numpy())
+
+        # Calculate cosine similarity between adjacent embeddings
+        coherence_scores = []
+        for i in range(len(embeddings) - 1):
+            sim = cosine_similarity(embeddings[i], embeddings[i + 1])[0][0]
+            coherence_scores.append(sim)
+
+        # Average coherence score
+        avg_coherence = (
+            sum(coherence_scores) / len(coherence_scores) if coherence_scores else 0.0
+        )
+        return avg_coherence
+
+
+class FluencyScorer(StoryScorer):
+    def __init__(
+        self,
+        model: PreTrainedModel,
+        tokenizer: PreTrainedTokenizer,
+        device: torch.device,
+    ):
+        self.model = model
+        self.tokenizer = tokenizer
+        self.device = device
+
+    def score(self, story: str) -> float:
+        # Mask each token in the story and calculate the probability of the original token
+        # Fluency is measured by the average probability of each token in the story
+        inputs = self.tokenizer(story, return_tensors="pt").to(self.device)
+        input_ids = inputs.input_ids
+        mask_token_id = self.tokenizer.mask_token_id
+
+        if mask_token_id is None:
+            self.tokenizer.mask_token = "[MASK]"
+            mask_token_id = self.tokenizer.encode(self.tokenizer.mask_token)[0]
+
+        fluency_scores = []
+        for i in range(1, input_ids.size(1) - 1):
+            masked_input_ids = input_ids.clone()
+            masked_input_ids[0, i] = mask_token_id
+
+            with torch.no_grad():
+                outputs = self.model(input_ids=masked_input_ids)
+                logits = outputs.logits
+
+            original_token_id = input_ids[0, i]
+            token_probability = logits[0, i].softmax(dim=-1)[original_token_id].item()
+            fluency_scores.append(token_probability)
+
+        avg_fluency = (
+            sum(fluency_scores) / len(fluency_scores) if fluency_scores else 0.0
+        )
+        return avg_fluency
+
+
+class GenreAlignmentScorer(StoryScorer):
+    def __init__(self, pipeline: Pipeline, genre: str):
+        self.pipeline = pipeline
+        self.genre = genre
+
+    def score(self, story: str) -> float:
+        if not self.genre:
+            return 0.5
+
+        # Evaluate by sentence to check whether the genre is maintained throughout
+        sentences = [s.strip() for s in story.split(".") if s.strip()]
+        results = []
+        for sentence in sentences:
+            result = self.pipeline(
+                sentence, candidate_labels=[self.genre], multi_label=True
+            )
+            results.append(result["scores"][0])
+
+        avg_core = sum(results) / len(results) if results else 0.0
+        return avg_core
+
+
+class StoryEvaluator:
+    def __init__(
+        self,
+        coherence_scorer: CoherenceScorer,
+        fluency_scorer: FluencyScorer,
+        genre_scorer: GenreAlignmentScorer,
+        weights: tuple[float, float, float] = (0.4, 0.3, 0.3),
+    ):
+        self.coherence_scorer = coherence_scorer
+        self.fluency_scorer = fluency_scorer
+        self.genre_scorer = genre_scorer
+        self.weights = weights
+
+    def evaluate(self, story: str, max_scores: list[float]) -> CombinedScore:
+        coherence = self.coherence_scorer.score(story)
+        fluency = self.fluency_scorer.score(story)
+        genre_alignment = self.genre_scorer.score(story)
+
+        max_scores[0] = np.max([max_scores[0], coherence])
+        max_scores[1] = np.max([max_scores[1], fluency])
+        max_scores[2] = np.max([max_scores[2], genre_alignment])
+
+        return CombinedScore(
+            coherence=coherence,
+            fluency=fluency,
+            genre_alignment=genre_alignment,
+            total=0,
+        )
+
+    def evaluate_multiple(self, stories: list[str]) -> list[tuple[str, CombinedScore]]:
+        """Evaluate multiple stories and return them sorted by total score."""
+
+        # Scores are normalized by the max scores on every evaluation
+        # This is to ensure that the scores are comparable between each other, as they are originally on different scales
+        
+        # Reset max scores
+        max_scores = [0.0, 0.0, 0.0]
+
+        scored_stories = [
+            (story, self.evaluate(story, max_scores)) for story in stories
+        ]
+
+        # Normalize scores
+        for _, scores in scored_stories:
+            scores.coherence, scores.fluency, scores.genre_alignment = np.divide(
+                [scores.coherence, scores.fluency, scores.genre_alignment],
+                max_scores,
+            )
+            scores.total = np.dot(
+                [scores.coherence, scores.fluency, scores.genre_alignment], self.weights
+            )
+
+        return sorted(scored_stories, key=lambda x: x[1].total, reverse=True)

story_beam_search/stories_generator.py

@@ -0,0 +1,191 @@
+import os
+import re
+from dataclasses import dataclass
+from story_beam_search.scoring import CombinedScore
+from transformers import (
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    AutoModelForMaskedLM,
+    pipeline,
+    Pipeline,
+)
+import torch
+
+auth_token = os.getenv("HF_TOKEN", None)
+
+
+@dataclass
+class ModelConfig:
+    text_model_name: str = "meta-llama/Llama-3.2-1B-Instruct"
+    bert_name: str = "bert-base-uncased"  # "answerdotai/ModernBERT-base"
+    zero_shot_name: str = "facebook/bart-large-mnli"
+    device: str = (
+        "mps"
+        if torch.backends.mps.is_available()
+        else "cuda" if torch.cuda.is_available() else "cpu"
+    )
+
+
+@dataclass
+class Models:
+    """Container for all loaded models and tokenizers."""
+
+    device: torch.device
+    text_model: AutoModelForCausalLM
+    text_tokenizer: AutoTokenizer
+    bert_model: AutoModelForMaskedLM
+    bert_tokenizer: AutoTokenizer
+    zero_shot_pipeline: Pipeline
+
+
+class ModelLoader:
+    """Handles loading and initialization of all required models."""
+
+    def __init__(self, config: ModelConfig = ModelConfig()):
+        self.config = config
+        self.device = torch.device(config.device)
+
+    def load_models(self) -> Models:
+        """Load all required models and return them in a Models container."""
+
+        # Load Text model for writting stories
+        print(f"Loading Text model ({self.config.text_model_name})...")
+        text_tokenizer = AutoTokenizer.from_pretrained(
+            self.config.text_model_name, token=auth_token
+        )
+        text_model = AutoModelForCausalLM.from_pretrained(
+            self.config.text_model_name
+        ).to(self.device)
+        text_model.eval()
+
+        # Load BERT model for coherence and fluency scoring
+        print(f"Loading BERT model ({self.config.bert_name})...")
+        bert_tokenizer = AutoTokenizer.from_pretrained(self.config.bert_name)
+        bert_model = AutoModelForMaskedLM.from_pretrained(self.config.bert_name).to(
+            self.device
+        )
+        bert_model.eval()
+
+        # Load Zero-Shot classification pipeline for genre alignment scoring
+        print("Loading Zero-Shot Classification pipeline...")
+        zero_shot_pipeline = pipeline(
+            "zero-shot-classification",
+            model=self.config.zero_shot_name,
+            device=self.device,
+        )
+
+        return Models(
+            device=self.device,
+            text_model=text_model,
+            text_tokenizer=text_tokenizer,
+            bert_model=bert_model,
+            bert_tokenizer=bert_tokenizer,
+            zero_shot_pipeline=zero_shot_pipeline,
+        )
+
+
+class StoryGenerationSystem:
+    """
+    High-level class that coordinates model loading and initialization of all components.
+    Acts as a facade for the entire story generation system.
+    """
+
+    def __init__(self, model_config: ModelConfig = ModelConfig()):
+        self.model_loader = ModelLoader(model_config)
+        self.models = None
+        self.beam_search = None
+        self.evaluator = None
+        self.storyness = None
+        self.injection_guard = None
+
+    def initialize(self) -> None:
+        """Initialize all components of the story generation system."""
+        from story_beam_search.beam_search import BeamSearchGenerator, BeamSearchConfig
+        from story_beam_search.scoring import (
+            CoherenceScorer,
+            FluencyScorer,
+            GenreAlignmentScorer,
+            StoryEvaluator,
+        )
+
+        # Load all models
+        self.models = self.model_loader.load_models()
+
+        # Initialize beam search
+        self.beam_search = BeamSearchGenerator(
+            model=self.models.text_model,
+            tokenizer=self.models.text_tokenizer,
+            device=self.models.device,
+            config=BeamSearchConfig(),
+        )
+
+        # Initialize scorers
+        coherence_scorer = CoherenceScorer(
+            model=self.models.bert_model,
+            tokenizer=self.models.bert_tokenizer,
+            device=self.models.device,
+        )
+
+        fluency_scorer = FluencyScorer(
+            model=self.models.text_model,
+            tokenizer=self.models.text_tokenizer,
+            device=self.models.device,
+        )
+
+        # Note: genre_scorer will be created per request as it depends on the user's genre choice
+        self.create_evaluator = lambda genre: StoryEvaluator(
+            coherence_scorer=coherence_scorer,
+            fluency_scorer=fluency_scorer,
+            genre_scorer=GenreAlignmentScorer(
+                pipeline=self.models.zero_shot_pipeline, genre=genre
+            ),
+        )
+
+        # Last minute addition 'misusing' the GenreAlignmentScorer to check for prompt injections
+        self.storyness = GenreAlignmentScorer(
+            pipeline=self.models.zero_shot_pipeline, genre="story"
+        )
+        self.injection_guard = GenreAlignmentScorer(
+            pipeline=self.models.zero_shot_pipeline, genre="prompt injection"
+        )
+
+    def generate_and_evaluate(
+        self, prompt: str, genre: str, num_stories: int = 3
+    ) -> list[tuple[str, CombinedScore]]:
+        """Generate stories and evaluate them."""
+        if not self.models:
+            raise RuntimeError("System not initialized. Call initialize() first.")
+        
+        # Low effort attempt to detect prompt injections using the zero-shot classifier
+        prompt_segments = re.split(r'[^a-zA-Z0-9 ]+', prompt)
+        prompt_segments = list(set(prompt_segments))
+
+        storyness_score = self.storyness.score(prompt)
+        for segment in prompt_segments:
+            if segment.strip():    
+                injection_score = self.injection_guard.score(segment)
+                if storyness_score < 0.2 or injection_score > 0.2:
+                    print("Potential prompt injection detected.")
+                    print(f"storyness_score: {storyness_score}")
+                    print(f"injection_score: {injection_score}")
+                    print("Prompt:", segment)
+                    raise ValueError("Prompt does not seem like a story. Please try again.")
+                 
+
+        # Create evaluator with specified genre
+        evaluator = self.create_evaluator(genre)
+
+        # Generate stories
+        # This is not strict beam search, inspired by https://huggingface.co/spaces/HuggingFaceH4/blogpost-scaling-test-time-compute
+        # to generate more diverse stories
+        all_stories = []
+        for _ in range(num_stories):
+            stories = self.beam_search.generate_iterations(prompt, genre, evaluator)
+            ranked_stories = evaluator.evaluate_multiple(stories)
+            # keep the top story of this beam search iteration
+            all_stories.append(ranked_stories[0][0])
+
+        # Evaluate stories once more
+        ranked_stories = evaluator.evaluate_multiple(all_stories)
+        # Return top k stories with their scores
+        return ranked_stories[:num_stories]