Image2Story / src /models /storyGeneration.py
ChinarQ-AI's picture
Create models//storyGeneration.py (#2)
0c31a3b verified
raw
history blame
413 Bytes
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)