src/models/storyGeneration.py

from langchain_huggingface import HuggingFaceEndpoint
import os

hf_token = os.environ.get("HUGGINGFACE_API_TOKEN")

story_generator = HuggingFaceEndpoint(
    repo_id="mistralai/Mistral-7B-Instruct-v0.3",
    huggingfacehub_api_token=hf_token,
    temperature=0.7,
    max_new_tokens=200
)

def generateStory(prompt):
    if not prompt:
        return "Invalid caption input."
    return story_generator.invoke(prompt)