Update app.py

app.py

@@ -55,9 +55,7 @@ st.header("Chat with your document 📄  (Model: Falcon-7B-Instruct)")
 model_name = "tiiuae/falcon-7b-instruct"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 
-   
-
-class CustomHuggingFaceEndpoint(HuggingFaceEndpoint):
+class CustomHuggingFaceEndpoint(HuggingFaceHub):
     def transform_input(self, prompt: str, model_kwargs: Dict) -> bytes:
         len_prompt = len(prompt)
         input_str = json.dumps({
@@ -78,8 +76,9 @@ class CustomHuggingFaceEndpoint(HuggingFaceEndpoint):
         ans = ans[:ans.rfind("Human")].strip()
         return ans
 
+
 def load_chain():
-    llm = CustomHuggingFaceEndpoint(model_name=model_name)
+    llm = CustomHuggingFaceEndpoint(repo_id=model_name)
     memory = ConversationBufferMemory()
     chain = ConversationChain(llm=llm, memory=memory)
     return chain