Create models//storyGeneration.py (#2)

- Create models//storyGeneration.py (c0c042bdb0f56aec179f3b85a0c87b5061e8b2a9)


Co-authored-by: Khan <[email protected]>

src/models/storyGeneration.py

@@ -0,0 +1,16 @@
+from langchain_huggingface import HuggingFaceEndpoint
+import os
+
+hf_token = os.environ.get("HUGGINGFACE_API_TOKEN")
+
+story_generator = HuggingFaceEndpoint(
+    repo_id="HuggingFaceH4/zephyr-7b-beta",
+    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)