Update app.py

app.py

@@ -22,16 +22,22 @@ def generate_caption_from_image(image_path):
 
 def generate_story_from_caption(caption):
     # Generate story based on caption
-    api_key = os.getenv("GOOGLE_API")
-    prompt_template = """You are a story teller;
-    You can generate a short story based on a simple narrative, the story should between 30 to 80 words;
-    CONTEXT: {scenario}
-    Story: """
-    PROMPT = PromptTemplate(template=prompt_template, input_variables=["scenario"])
-    llm_chain = LLMChain(prompt=PROMPT,
-                        llm=GooglePalm(google_api_key=api_key, temperature=0.8))
-    scenario = caption
-    story = llm_chain.run(scenario)
+    llm = HuggingFaceHub(huggingfacehub_api_token = os.getenv('HUGGING_FACE')}"},
+                        repo_id = "tiiuae/falcon-7b-instruct",
+                        verbose = False,
+                        model_kwargs = {"temperature":0.2, "max_new_tokens": 4000})
+    template= """You are a story teller.
+    You get a scenario as an input text, and generates a short story out of it.
+    Context: {scenario}
+    Story:
+    """
+    prompt = PromptTemplate(template=template, input_variables=["scenario"])
+    #Let's create our LLM chain now
+    chain = LLMChain(prompt=prompt, llm=llm)
+    story = chain.run(caption)
+    start_index = story.find("Story:") + len("Story:")
+    # Extract the text after "Story:"
+    story = story[start_index:].strip()
     return story
 
 def text_to_speech(text):