Update space with latest code and dependencies on Thu Jan 2 21:22:45 UTC 2025

README.md

@@ -9,8 +9,8 @@ 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
+- openai-community/gpt2
+- answerdotai/ModernBERT-base
 - facebook/bart-large-mnli
 license: mit
 ---

app.py

@@ -16,6 +16,7 @@ genre_choices = [
     "horror",
 ]
 
+
 class InputModel(BaseModel):
     prompt: str
     genre: str
@@ -28,7 +29,7 @@ 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]]:
@@ -39,7 +40,11 @@ def create_story_generation_interface() -> gr.Interface:
 
         # 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
+            prompt=prompt,
+            genre=genre,
+            num_stories=num_stories,
+            temperature=temperature,
+            max_length=max_length,
         )
 
         # Update beam search config with user parameters
@@ -48,7 +53,9 @@ def create_story_generation_interface() -> gr.Interface:
 
         # Generate and evaluate stories
         ranked_stories = system.generate_and_evaluate(
-            input_values.prompt, input_values.genre, num_stories=input_values.num_stories
+            input_values.prompt,
+            input_values.genre,
+            num_stories=input_values.num_stories,
         )
 
         # Format detailed scores

requirements.txt

@@ -1,3 +1,4 @@
+accelerate==1.2.1 ; python_full_version == "3.10.13"
 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"
@@ -30,6 +31,7 @@ packaging==24.2 ; python_full_version == "3.10.13"
 pandas==2.2.3 ; python_full_version == "3.10.13"
 pillow==11.1.0 ; python_full_version == "3.10.13"
 protobuf==5.29.2 ; python_full_version == "3.10.13"
+psutil==6.1.1 ; 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"
@@ -41,9 +43,9 @@ 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"
+ruff==0.8.5 ; 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"
+safetensors==0.5.0 ; 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"
@@ -57,7 +59,7 @@ 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"
+transformers @ git+https://github.com/huggingface/transformers.git@42865860ec6dc135972d9555753cb7ee17f51fb4 ; 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"

story_beam_search/beam_search.py

@@ -51,7 +51,7 @@ class BeamSearchGenerator:
             [story[instructions_len:] for story in stories]
         )
 
-        stories = [story for story, _ in ranked_stories[:self.config.num_beams]]
+        stories = [story for story, _ in ranked_stories[: self.config.num_beams]]
 
         if stories:
             for _ in range(self.config.num_iterations):
@@ -64,7 +64,9 @@ class BeamSearchGenerator:
                 ranked_stories = evaluator.evaluate_multiple(
                     [story[instructions_len:] for story in all_stories]
                 )
-                stories = [story for story, _ in ranked_stories[:self.config.num_beams]]
+                stories = [
+                    story for story, _ in ranked_stories[: self.config.num_beams]
+                ]
 
         return stories
 

story_beam_search/scoring.py

@@ -46,7 +46,9 @@ class CoherenceScorer(StoryScorer):
         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, :]
+                outputs = self.model(**inputs)
+                last_hidden_state = outputs.hidden_states[-1]
+                emb = last_hidden_state[:, 0, :]
                 embeddings.append(emb.cpu().numpy())
 
         # Calculate cosine similarity between adjacent embeddings

story_beam_search/stories_generator.py

@@ -16,8 +16,8 @@ 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"
+    text_model_name: str = "openai-community/gpt2"
+    bert_name: str = "answerdotai/ModernBERT-base"
     zero_shot_name: str = "facebook/bart-large-mnli"
     device: str = (
         "mps"
@@ -54,14 +54,14 @@ class ModelLoader:
             self.config.text_model_name, token=auth_token
         )
         text_model = AutoModelForCausalLM.from_pretrained(
-            self.config.text_model_name
+            self.config.text_model_name, device_map="auto", torch_dtype=torch.float16
         ).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(
+        bert_model = AutoModelForMaskedLM.from_pretrained(self.config.bert_name, output_hidden_states=True).to(
             self.device
         )
         bert_model.eval()
@@ -155,22 +155,23 @@ class StoryGenerationSystem:
         """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 = 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():    
+            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.")
-                 
+                    raise ValueError(
+                        "Prompt does not seem like a story. Please try again."
+                    )
 
         # Create evaluator with specified genre
         evaluator = self.create_evaluator(genre)