Update components/custom_llm.py

components/custom_llm.py

@@ -6,6 +6,22 @@ from typing import Literal
 import requests
 from langchain.prompts import PromptTemplate, ChatPromptTemplate
 from operator import itemgetter
+import re
+
+
+def format_captions(text):
+  sentences = list(filter(None,[x.strip() for x in re.split(r'[^A-Za-z0-9 -]', text)]))
+  print(len(sentences))
+  return "\n".join([f"{i+1}. {d}" for i,d in enumerate(sentences)])
+
+def custom_chain():
+
+    prompt = PromptTemplate.from_template("<s><INST>Given the below template, create a list of image generation prompt with maximum 5 words for each number\n\n{template}<INST> ")
+    
+    cap_llm = CustomLLM(repo_id="mistralai/Mixtral-8x7B-Instruct-v0.1", model_type='text-generation', api_token=API_TOKEN, stop=["\n<|"])
+    
+    return {"template":lambda x:format_captions(x)} | prompt | cap_llm
+
 
 class CustomLLM(LLM):
     repo_id : str