Update story_generator.py

story_generator.py

@@ -1,119 +1,189 @@
-import os
 import torch
-from fastapi import FastAPI, HTTPException
-from fastapi.responses import JSONResponse
-from pydantic import BaseModel
-import story_generator
-from diffusers import DiffusionPipeline
-from PIL import Image
-import io
-import base64
-import time
+from transformers import GPT2Tokenizer, GPT2LMHeadModel
+import random
+import re
 
-app = FastAPI()
+# Set Hugging Face cache directory
+CACHE_DIR = "/tmp/huggingface"
 
-# Set Hugging Face cache directories
-os.environ["HF_HOME"] = "/tmp/huggingface"
-os.environ["TRANSFORMERS_CACHE"] = "/tmp/huggingface"
-os.environ["HF_HUB_CACHE"] = "/tmp/huggingface"
+# ------------------------
+# Load Story Generation Model
+# ------------------------
+STORY_MODEL_NAME = "abdalraheemdmd/story-api"
+story_tokenizer = GPT2Tokenizer.from_pretrained(STORY_MODEL_NAME, cache_dir=CACHE_DIR)
+story_model = GPT2LMHeadModel.from_pretrained(STORY_MODEL_NAME, cache_dir=CACHE_DIR)
 
-# Enable GPU if available
-device = "cuda" if torch.cuda.is_available() else "cpu"
+# ------------------------
+# Load Question Generation Model
+# ------------------------
+QUESTION_MODEL_NAME = "abdalraheemdmd/question-gene"
+question_tokenizer = GPT2Tokenizer.from_pretrained(QUESTION_MODEL_NAME, cache_dir=CACHE_DIR)
+question_model = GPT2LMHeadModel.from_pretrained(QUESTION_MODEL_NAME, cache_dir=CACHE_DIR)
 
-# Load image generation model
-IMAGE_MODEL = "lykon/dreamshaper-8"
-pipeline = DiffusionPipeline.from_pretrained(
-    IMAGE_MODEL,
-    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
-).to(device)
+# Ensure tokenizers have a pad token
+if story_tokenizer.pad_token_id is None:
+    story_tokenizer.pad_token_id = story_tokenizer.eos_token_id
+if question_tokenizer.pad_token_id is None:
+    question_tokenizer.pad_token_id = question_tokenizer.eos_token_id
 
-# Define request schema
-class StoryRequest(BaseModel):
-    theme: str
-    reading_level: str
+def generate_story(theme, reading_level, max_new_tokens=400, temperature=0.7):
+    """Generates a story based on the provided theme and reading level."""
+    prompt = f"A {reading_level} story about {theme}:"
+    input_ids = story_tokenizer(prompt, return_tensors="pt").input_ids
+    with torch.no_grad():
+        output = story_model.generate(
+            input_ids,
+            max_new_tokens=max_new_tokens,
+            temperature=temperature,
+            top_k=20,
+            top_p=0.7,
+            do_sample=True,
+            early_stopping=True,
+            pad_token_id=story_tokenizer.pad_token_id,
+            eos_token_id=story_tokenizer.eos_token_id,
+            attention_mask=input_ids.ne(story_tokenizer.pad_token_id)
+        )
+    return story_tokenizer.decode(output[0], skip_special_tokens=True)
 
-def generate_images(prompts, max_retries=3, delay=2):
+def extract_protagonist(story):
     """
-    Generates images in batch with a retry mechanism.
-    If a specific indexing error occurs, it retries up to max_retries times.
-    If all retries fail, it falls back to sequential generation.
+    Attempts to extract the protagonist from the first sentence by searching for the pattern "named <Name>".
+    Returns the first matched name, if available.
     """
-    for attempt in range(max_retries):
-        try:
-            print(f"Batch image generation attempt {attempt+1}...")
-            results = pipeline(prompt=prompts, num_inference_steps=15, height=768, width=768).images
-            return results
-        except Exception as e:
-            if "index 16 is out of bounds" in str(e):
-                print(f"Encountered indexing error on attempt {attempt+1}: {e}")
-                time.sleep(delay)  # wait before retrying
-            else:
-                raise e
-    # Fallback: Sequential generation
-    print("Falling back to sequential image generation...")
-    images = []
-    for i, prompt in enumerate(prompts):
-        try:
-            print(f"Generating image for prompt {i+1} sequentially...")
-            image = pipeline(prompt=prompt, num_inference_steps=15, height=768, width=768).images[0]
-            images.append(image)
-        except Exception as e:
-            print(f"Sequential generation failed for prompt {i+1}: {e}")
-            raise e
-    return images
+    sentences = re.split(r'\.|\n', story)
+    if sentences:
+        m = re.search(r"named\s+([A-Z][a-z]+)", sentences[0])
+        if m:
+            return m.group(1)
+    return None
 
-@app.post("/generate_story_questions_images")
-def generate_story_questions_images(request: StoryRequest):
-    """Generates a story, dynamic questions, and cartoonish storybook images."""
-    try:
-        print(f"🎭 Generating story for theme: {request.theme} and level: {request.reading_level}")
-        
-        # Generate story and questions
-        story_result = story_generator.generate_story_and_questions(request.theme, request.reading_level)
-        story_text = story_result.get("story", "").strip()
-        questions = story_result.get("questions", "").strip()
-        
-        if not story_text:
-            raise HTTPException(status_code=500, detail="Story generation failed.")
-        
-        # Split the story into up to 6 paragraphs
-        paragraphs = [p.strip() for p in story_text.split("\n") if p.strip()][:6]
-        
-        # Build a list of prompts for batched image generation
-        prompts = [
-            (
-                f"Children's storybook illustration of: {p}. "
-                "Soft pastel colors, hand-drawn style, friendly characters, warm lighting, "
-                "fantasy setting, watercolor texture, storybook illustration, beautiful composition."
-            )
-            for p in paragraphs
-        ]
-        print(f"Generating images for {len(prompts)} paragraphs concurrently...")
-        
-        # Use the retry mechanism for image generation
-        results = generate_images(prompts, max_retries=3, delay=2)
-        
-        # Convert each image to Base64
-        images = []
-        for image in results:
-            img_byte_arr = io.BytesIO()
-            image.save(img_byte_arr, format="PNG")
-            img_byte_arr.seek(0)
-            base64_image = base64.b64encode(img_byte_arr.getvalue()).decode("utf-8")
-            images.append(base64_image)
-        
-        return JSONResponse(content={
-            "theme": request.theme,
-            "reading_level": request.reading_level,
-            "story": story_text,
-            "questions": questions,
-            "images": images
-        })
+def extract_characters(story):
+    """
+    Extracts potential character names from the story using a frequency count on capitalized words.
+    Filters out common stopwords so that the most frequently mentioned name is likely the main character.
+    """
+    words = re.findall(r'\b[A-Z][a-zA-Z]+\b', story)
+    stopwords = {"The", "A", "An", "And", "But", "Suddenly", "Quickly", "However", "Well",
+                 "They", "I", "He", "She", "It", "When", "Where", "Dr", "Mr"}
+    filtered = [w for w in words if w not in stopwords and len(w) > 2]
+    if not filtered:
+        return []
+    freq = {}
+    for word in filtered:
+        freq[word] = freq.get(word, 0) + 1
+    sorted_chars = sorted(freq.items(), key=lambda x: x[1], reverse=True)
+    return [item[0] for item in sorted_chars]
+
+def extract_themes(story):
+    """Extracts themes from the story based on keyword matching."""
+    themes = []
+    story_lower = story.lower()
+    if "space" in story_lower:
+        themes.append("space")
+    if "adventure" in story_lower:
+        themes.append("adventure")
+    if "friend" in story_lower:
+        themes.append("friendship")
+    if "learn" in story_lower or "lesson" in story_lower:
+        themes.append("learning")
+    return themes
+
+def extract_lesson(story):
+    """
+    Attempts to extract a lesson or moral from the story by finding sentences
+    containing keywords like "learn" or "lesson". Returns the last matching sentence.
+    """
+    sentences = re.split(r'\.|\n', story)
+    lesson_sentences = [
+        s.strip() for s in sentences
+        if ("learn" in s.lower() or "lesson" in s.lower()) and len(s.strip()) > 20
+    ]
+    if lesson_sentences:
+        return lesson_sentences[-1]
+    else:
+        return "No explicit lesson found."
+
+def format_question(question_prompt, correct_answer, distractors):
+    """
+    Combines the correct answer with three distractors, shuffles the options,
+    and formats the question as a multiple-choice question.
+    """
+    # Ensure exactly 3 distractors are available
+    if len(distractors) < 3:
+        default_distractors = ["Option X", "Option Y", "Option Z"]
+        while len(distractors) < 3:
+            distractors.append(default_distractors[len(distractors) % len(default_distractors)])
+    else:
+        distractors = random.sample(distractors, 3)
+    options = distractors + [correct_answer]
+    random.shuffle(options)
+    letters = ["A", "B", "C", "D"]
+    correct_letter = letters[options.index(correct_answer)]
+    options_text = "\n".join(f"{letters[i]}) {option}" for i, option in enumerate(options))
+    question_text = f"{question_prompt}\n{options_text}\nCorrect Answer: {correct_letter}"
+    return question_text
+
+def dynamic_fallback_questions(story):
+    """
+    Generates three multiple-choice questions based on dynamic story content.
+    Each question uses a randomly chosen template and shuffles its options.
+    """
+    protagonist = extract_protagonist(story)
+    characters = extract_characters(story)
+    themes = extract_themes(story)
+    lesson = extract_lesson(story)
+    
+    # --- Question 1: Theme ---
+    theme_templates = [
+        "What is the main theme of the story?",
+        "Which theme best represents the narrative?",
+        "What subject is central to the story?"
+    ]
+    q1_prompt = random.choice(theme_templates)
+    correct_theme = " and ".join(themes) if themes else "learning"
+    q1_distractors = ["sports and competition", "cooking and baking", "weather and seasons", "technology and innovation"]
+    q1 = format_question(q1_prompt, correct_theme, q1_distractors)
+    
+    # --- Question 2: Primary Character ---
+    character_templates = [
+        "Who is the primary character in the story?",
+        "Which character drives the main action in the narrative?",
+        "Who is the central figure in the story?"
+    ]
+    q2_prompt = random.choice(character_templates)
+    if protagonist:
+        correct_character = protagonist
+    elif characters:
+        correct_character = characters[0]
+    else:
+        correct_character = "The main character"
+    q2_distractors = ["a mysterious stranger", "an unknown visitor", "a supporting character", "a sidekick"]
+    q2 = format_question(q2_prompt, correct_character, q2_distractors)
     
-    except Exception as e:
-        print(f"❌ Error generating story/questions/images: {e}")
-        raise HTTPException(status_code=500, detail=str(e))
+    # --- Question 3: Lesson/Moral ---
+    lesson_templates = [
+        "What lesson did the characters learn by the end of the story?",
+        "What moral can be inferred from the narrative?",
+        "What is the key takeaway from the story?"
+    ]
+    q3_prompt = random.choice(lesson_templates)
+    if lesson and lesson != "No explicit lesson found.":
+        correct_lesson = lesson  # full sentence without truncation
+    else:
+        correct_lesson = "understanding and growth"
+    q3_distractors = ["always be silent", "never try new things", "do nothing", "ignore opportunities"]
+    q3 = format_question(q3_prompt, correct_lesson, q3_distractors)
+    
+    return f"{q1}\n\n{q2}\n\n{q3}"
+
+def generate_story_and_questions(theme, reading_level):
+    """
+    Generates a story using the story generation model and then creates dynamic,
+    multiple-choice questions based on that story.
+    """
+    story = generate_story(theme, reading_level)
+    questions = dynamic_fallback_questions(story)
+    return {"story": story, "questions": questions}
 
-@app.get("/")
-def home():
-    return {"message": "🎉 Welcome to the Story, Question & Image API!"}
+# Alias for backward compatibility
+create_fallback_questions = dynamic_fallback_questions