Upload brain.py

brain.py

@@ -0,0 +1,52 @@
+from dotenv import load_dotenv, find_dotenv
+from transformers import pipeline
+from transformers import pipeline
+from langchain_huggingface import HuggingFaceEndpoint
+from langchain_core.prompts import PromptTemplate
+load_dotenv(find_dotenv())
+
+
+def img2txt(image_path):
+    """ Convert image to text using Hugging Face pipeline.
+    Args:
+        image_path (str): Path to the image.
+    Returns:
+        str: The text extracted from the image.
+    """
+    itt = pipeline(
+        "image-to-text",
+        model="Salesforce/blip-image-captioning-base"
+    )
+    text = itt(image_path)[0]["generated_text"]
+    print(text)
+    return text
+
+
+def generate_story(scenario, repo_id="mistralai/Mistral-7B-Instruct-v0.2"):
+    """ Generate a story using image captioning and language model.
+    Args:
+        scenario (str): The scenario extracted from the image.
+    Returns:
+        str: The story generated using the scenario.
+    """
+    llm = HuggingFaceEndpoint(
+        repo_id=repo_id,
+        temperature=0.5,
+        streaming=True
+    )
+    prompt_template = """
+    You are a kids story writer. Provide a coherent story for kids
+    using this simple instruction: {scenario}. The story should have a clear
+    beginning, middle, and end. The story should be interesting and engaging for
+    kids. The story should be maximum 200 words long. Do not include
+    any adult or polemic content.
+    Story:
+    """
+    prompt = PromptTemplate.from_template(prompt_template)
+    story = prompt | llm
+    return story.invoke(input={"segmentation_results": scenario})
+
+
+if __name__ == "__main__":
+    my_story = generate_story(img2txt("image.jpg"))
+    print(my_story)