Create storyGeneration.py

src/models/storyGeneration.py

@@ -0,0 +1,18 @@
+from langchain_huggingface import HuggingFaceEndpoint
+import os
+from dotenv import load_dotenv
+
+load_dotenv()
+hf_token = os.getenv("HUGGINGFACE_API_TOKEN")
+
+story_generator = HuggingFaceEndpoint(
+    repo_id="HuggingFaceH4/zephyr-7b-beta",
+    huggingfacehub_api_token=hf_token,
+    temperature=0.7,
+    max_new_tokens=800
+)
+
+def generateStory(prompt):
+    if not prompt:
+        return "Invalid caption input."
+    return story_generator.invoke(prompt)