Upload 34 files

.gitignore

@@ -0,0 +1,162 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+#poetry.lock
+
+# pdm
+#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#pdm.lock
+#   pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
+#   in version control.
+#   https://pdm.fming.dev/latest/usage/project/#working-with-version-control
+.pdm.toml
+.pdm-python
+.pdm-build/
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+#  and can be added to the global gitignore or merged into this file.  For a more nuclear
+#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
+#.idea/

app/api/routes.py

@@ -0,0 +1,39 @@
+from fastapi import APIRouter, HTTPException
+from app.models.request_models import StoryRequest, StoryResponse
+from app.services.story_generator import StoryGenerator
+from app.services.story_refinement import clean_text, refine_story
+import logging
+
+# Create the router instance
+router = APIRouter()
+logger = logging.getLogger(__name__)
+
+@router.post("/generate_story", response_model=StoryResponse)
+async def generate_story(request: StoryRequest):
+    try:
+        # Initialize generator and generate initial story
+        generator = StoryGenerator()
+        initial_story = generator.generate_story(
+            request.prompt,
+            request.max_length,
+            request.temperature
+        )
+        
+        # Clean the initial story
+        cleaned_story = clean_text(initial_story)
+        
+        try:
+            # Try to get the refined story
+            final_story = await refine_story(cleaned_story)
+        except Exception as e:
+            logger.error(f"Refinement failed: {str(e)}")
+            final_story = cleaned_story  # Use cleaned story if refinement fails
+        
+        return StoryResponse(story=final_story)
+        
+    except Exception as e:
+        logger.error(f"Story generation failed: {str(e)}")
+        raise HTTPException(
+            status_code=500,
+            detail=f"Story generation failed: {str(e)}"
+        )

app/core/config.py

@@ -0,0 +1,14 @@
+import os
+from dotenv import load_dotenv
+
+load_dotenv()
+
+class Settings:
+    PROJECT_NAME: str = "Story Generator API"
+    MODEL_PATH: str = "best_model.pt"
+    GOOGLE_API_KEY: str = os.getenv('GEMINI_API_KEY')
+    
+    if not GOOGLE_API_KEY:
+        raise ValueError("Gemini API Key is not set in the environment variable GEMINI_API_KEY")
+
+settings = Settings()

app/models/request_models.py

@@ -0,0 +1,9 @@
+from pydantic import BaseModel, Field
+
+class StoryRequest(BaseModel):
+    prompt: str
+    max_length: int = Field(default=200, ge=1, le=1000)
+    temperature: float = Field(default=0.7, ge=0.1, le=1.0)
+
+class StoryResponse(BaseModel):
+    story: str

app/services/story_generator.py

@@ -0,0 +1,52 @@
+import torch
+from transformers import GPT2Tokenizer, GPT2LMHeadModel, GPT2Config
+from app.core.config import settings
+
+class StoryGenerator:
+    def __init__(self, model_path=settings.MODEL_PATH):
+        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+        print(f"Using device: {self.device}")
+        
+        self.tokenizer = GPT2Tokenizer.from_pretrained('gpt2')
+        self.tokenizer.pad_token = self.tokenizer.eos_token
+        
+        self.config = GPT2Config(
+            vocab_size=self.tokenizer.vocab_size,
+            n_positions=256,
+            n_ctx=256,
+            n_embd=256,
+            n_layer=4,
+            n_head=8
+        )
+        
+        self.model = GPT2LMHeadModel(self.config)
+        self.load_model(model_path)
+        self.model.to(self.device)
+        self.best_loss = float('inf')
+    
+    def load_model(self, path):
+        checkpoint = torch.load(path, map_location=self.device)
+        self.model.load_state_dict(checkpoint['model_state_dict'])
+
+    def generate_story(self, prompt, max_length=200, temperature=0.7):
+        self.model.eval()
+        input_ids = self.tokenizer.encode(prompt, return_tensors='pt').to(self.device)
+        
+        with torch.no_grad():
+            output_ids = self.model.generate(
+                input_ids,
+                max_length=max_length,
+                num_return_sequences=1,
+                do_sample=True,
+                top_k=50,
+                top_p=0.95,
+                num_beams=1,
+                temperature=temperature,
+                pad_token_id=self.tokenizer.pad_token_id,
+                eos_token_id=self.tokenizer.eos_token_id,
+                length_penalty=1.0,
+                no_repeat_ngram_size=3,
+                early_stopping=True
+            )[0]
+        
+        return self.tokenizer.decode(output_ids, skip_special_tokens=True)

app/services/story_refinement.py

@@ -0,0 +1,53 @@
+import re
+import unicodedata
+import google.generativeai as genai
+from fastapi import HTTPException
+from app.core.config import settings
+
+# Configure Gemini API
+genai.configure(api_key=settings.GOOGLE_API_KEY)
+gemini_model = genai.GenerativeModel('gemini-1.5-flash')
+
+def clean_text(story: str) -> str:
+    """Clean and normalize the generated text"""
+    cleaned_story = unicodedata.normalize("NFKD", story)
+    cleaned_story = re.sub(r'[\n\\]+', ' ', cleaned_story)
+    cleaned_story = re.sub(r'\s+', ' ', cleaned_story)
+    cleaned_story = (
+        cleaned_story.replace("“", "\"").replace("â€", "\"").replace("’", "'")
+                     .replace("â€TMs", "'s").replace("“", "\"").replace("Ó", "O")
+                     .replace("â", "").replace("�", "")
+    )
+    cleaned_story = cleaned_story.replace("\\", "")
+    return cleaned_story.strip()
+
+async def refine_story(story: str) -> str:
+    """Refine the story using the Gemini API with two-step refinement"""
+    try:
+        # First refinement
+        first_prompt = f"""
+        Refine the following story to make it coherent, grammatically correct, and fluent. 
+        Maintain the original plot and ending, but enhance readability, fluency, and consistency:
+
+        {story}
+        """
+        
+        first_response = gemini_model.generate_content(first_prompt)
+        first_refined = first_response.text if hasattr(first_response, 'text') else str(first_response)
+        
+        # Second refinement for final polishing
+        final_prompt = f"""
+        Please perform a final polish on this story to ensure perfect coherence, 
+        flow, and narrative structure while maintaining the essence of the story:
+
+        {first_refined}
+        """
+        
+        final_response = gemini_model.generate_content(final_prompt)
+        final_refined = final_response.text if hasattr(final_response, 'text') else str(final_response)
+        
+        return final_refined if final_refined else "No refined story generated."
+    
+    except Exception as e:
+        print(f"Error in story refinement: {str(e)}")
+        return story  # Return original story if refinement fails

best_model.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0e56be320b521adf85965adbfd542e981a492ce11870306f61d03a9f161e53b7
+size 64385826

main.py

@@ -0,0 +1,21 @@
+from fastapi import FastAPI
+from fastapi.middleware.cors import CORSMiddleware
+from app.api.routes import router
+
+app = FastAPI(title="Story Generator API")
+
+# Add CORS middleware
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["http://localhost:3000"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+# Include the router
+app.include_router(router)
+
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=8000)

requirements.txt

@@ -0,0 +1,7 @@
+fastapi==0.115.5
+protobuf==5.28.3
+pydantic==2.9.2
+python-dotenv==1.0.1
+torch==2.4.1
+transformers==4.44.2
+uvicorn==0.32.0