Update components/custom_llm.py

components/custom_llm.py

@@ -17,9 +17,13 @@ def format_captions(text):
 def custom_chain():
     API_TOKEN = os.environ['HF_INFER_API']
 
-    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> ")
+    # 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<|"])
+    # cap_llm = CustomLLM(repo_id="mistralai/Mixtral-8x7B-Instruct-v0.1", model_type='text-generation', api_token=API_TOKEN, stop=["\n<|"])
+    prompt = PromptTemplate.from_template("<s><INST>Given the below template, for each number, create a detailed description with maximum one word\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<|"], temperature=0.7)
+
     
     return {"template":lambda x:format_captions(x)} | prompt | cap_llm