updated app.py

app.py

@@ -1,7 +1,8 @@
 from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
+from langchain_community.llms import LlamaCpp
 from huggingface_hub.file_download import http_get
-from llama_cpp import Llama
+# from llama_cpp import Llama
 from langchain_core.callbacks import CallbackManager, StreamingStdOutCallbackHandler
 # from langchain_core.prompts import ChatPromptTemplate
 
@@ -19,7 +20,6 @@ class StoryRequest(BaseModel):
     mood: str
     story_type: str
     theme: str
-    length: int
     num_scenes: int
     txt: str
 
@@ -40,9 +40,19 @@ def load_model(
     os.chmod(final_model_path, 0o777)
     print("Files downloaded!")
     
-    model = Llama(
+    # model = Llama(
+    #     model_path=final_model_path,
+    #     n_ctx=1024
+    # )
+
+    model = LlamaCpp(
         model_path=final_model_path,
-        n_ctx=1024
+        temperature=0.3,
+        max_tokens=2000,
+        top_p=1,
+        n_ctx=1024,
+        callback_manager=callback_manager,
+        verbose=True,
     )
     
     print("Model loaded!")
@@ -51,6 +61,7 @@ def load_model(
 
 llm = load_model()
 
+
 # Create a prompt template
 # system = """You are a helpful and creative assistant that specializes in generating engaging and imaginative stories for kids.
 # Based on the user's provided mood, preferred story type, theme, age, and desired story length of 500-600 words, create a unique and captivating story.
@@ -58,71 +69,37 @@ llm = load_model()
 # to create another adventure soon! 
 # """
 
-# system = """You are a helpful and creative assistant that specializes in generating engaging and imaginative short storie for kids.
-# Based on the user's provided mood, preferred story type, theme, age, and desired story length of 500-600 words, create a unique and captivating story.
-# Always start with Story Title then generate a single story.Storie begin on Page 1(also mention the all pages headings in bold) and end on Page 7.
-# Total pages in storie are seven each page have one short paragraph and dont ask for any feedback at the end just sign off with a cute closing inviting the reader 
-# to create another adventure soon! 
-# """
-
-from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
-from langchain_core.messages import HumanMessage
-
-prompt = ChatPromptTemplate.from_messages(
-    [
-        (
-            "system",
-            """You are a helpful and creative assistant that specializes in generating engaging and imaginative short storie for kids.
-                Based on the user's provided mood, preferred story type, theme, age, and desired story length of 500-600 words, create a unique and captivating story.
-                Always start with Story Title then generate a single story.Storie begin on Page 1(also mention the all pages headings in bold) and end on Page 7.
-                Total pages in storie are seven each page have one short paragraph and dont ask for any feedback at the end just sign off with a cute closing inviting the reader 
-                to create another adventure soon! 
-                """,
-        ),
-        MessagesPlaceholder(variable_name="messages"),
-    ]
-)
-
-
+system = """You are a helpful and creative assistant that specializes in generating engaging and imaginative short storie for kids.
+Based on the user's provided mood, preferred story type, theme, age, and desired story length of 500-600 words, create a unique and captivating story.
+Always start with Story Title then generate a single story.Storie begin on Page 1(also mention the all pages headings in bold) and end on Page 7.
+Total pages in storie are seven each page have one short paragraph and dont ask for any feedback at the end just sign off with a cute closing inviting the reader 
+to create another adventure soon! 
+"""
 
-# prompt_template = ChatPromptTemplate.from_messages([("system", system), ("human", "{text}")])
+prompt_template = ChatPromptTemplate.from_messages([("system", system), ("human", "{text}")])
 
 # FastAPI endpoint to generate the story
 @app.post("/generate_story/")
 async def generate_story(story_request: StoryRequest):
-
-    # chain = prompt | llm
-    
     story = f"""here are the inputs from user:
     - **Mood:** {story_request.mood}
     - **Story Type:** {story_request.story_type}
     - **Theme:** {story_request.theme}
     - **Details Provided:** {story_request.txt}
     """
-
-    # response = chain.invoke({"messages": [HumanMessage(content=story)]})
-
-    formatted_prompt = prompt.format(messages=[HumanMessage(content=story)])
     
-    # Extract the raw string from the `ChatPromptValue`
-    # formatted_text = formatted_prompt.to_string()
-
-    chain = prompt | llm
-    response = chain.invoke(formatted_prompt)
-    # response = llm(formatted_prompt)
-    
-    # final_prompt = prompt_template.format(text=story)
+    final_prompt = prompt_template.format(text=story)
 
     # Create the LLMChain
     # chain = LLMChain(llm=llm, prompt=prompt_template)
-    # chain = llm | prompt_template
+    chain = llm | prompt_template
     
     # try:
     #     response = chain.invoke(final_prompt)
     #     return {"story": response}
     # except Exception as e:
     #     raise HTTPException(status_code=500, detail=str(e))
-    # response = chain.invoke(final_prompt)
+    response = chain.invoke(final_prompt)
     
     if not response:
         raise HTTPException(status_code=500, detail="Failed to generate the story")
@@ -152,14 +129,4 @@ async def generate_story(story_request: StoryRequest):
     return {
         "story": response,
         "images": images
-    }
-    
-
-
-# image_prompt = (
-#         f"Generate an image for Scene {i+1}. "
-#         f"This image should represent the details described in paragraph {i+1} of the story. "
-#         f"Mood: {mood}, Story Type: {', '.join(story_type)}, Theme: {theme}. "
-#         f"Story: {response} "
-#         f"Focus on the key elements in paragraph {i+1}."
-#     )
+    }