Update components/custom_llm.py

components/custom_llm.py

@@ -6,7 +6,7 @@ from typing import Literal
 import requests
 from langchain.prompts import PromptTemplate, ChatPromptTemplate
 from operator import itemgetter
-import re
+import re, os
 
 
 def format_captions(text):
@@ -15,6 +15,7 @@ def format_captions(text):
   return "\n".join([f"{i+1}. {d}" for i,d in enumerate(sentences)])
 
 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> ")