Update app.py

app.py

@@ -56,13 +56,10 @@ model_name = "tiiuae/falcon-7b-instruct"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 
    
-class ContentHandler(HuggingFaceEndpoint.CallbackHandler):
-    def __init__(self, model_name):
-        self.model_name = model_name
-        self.len_prompt = 0
 
+class CustomHuggingFaceEndpoint(HuggingFaceEndpoint):
     def transform_input(self, prompt: str, model_kwargs: Dict) -> bytes:
-        self.len_prompt = len(prompt)
+        len_prompt = len(prompt)
         input_str = json.dumps({
             "inputs": prompt,
             "parameters": {
@@ -73,7 +70,7 @@ class ContentHandler(HuggingFaceEndpoint.CallbackHandler):
             }
         })
         return input_str.encode('utf-8')
-         
+
     def transform_output(self, output: bytes) -> str:
         response_json = output.decode('utf-8')
         res = json.loads(response_json)
@@ -81,13 +78,8 @@ class ContentHandler(HuggingFaceEndpoint.CallbackHandler):
         ans = ans[:ans.rfind("Human")].strip()
         return ans
 
-content_handler = ContentHandler(model_name=model_name)
-
 def load_chain():
-    llm = HuggingFaceEndpoint(
-        model_name=model_name,
-        content_handler=content_handler,
-    )
+    llm = CustomHuggingFaceEndpoint(model_name=model_name)
     memory = ConversationBufferMemory()
     chain = ConversationChain(llm=llm, memory=memory)
     return chain