Update app.py

app.py

@@ -56,55 +56,12 @@ async def init():
     # docsearch = await cl.make_async(Chroma.from_documents)(pdf_data, embeddings) 
     docsearch = Chroma.from_documents(pdf_data, embeddings)
 
-    # custom SageMaker Model
-    class Llama2SageMaker(LLM):
-        max_new_tokens: int = 256
-        top_p: float = 0.9
-        temperature: float = 0.1
-    
-        @property
-        def _llm_type(self) -> str:
-            return "Llama2SageMaker"
-    
-        def _call(
-            self,
-            prompt: str,
-            stop: Optional[List[str]] = None,
-            run_manager: Optional[CallbackManagerForLLMRun] = None,
-        ) -> str:
-            if stop is not None:
-                raise ValueError("stop kwargs are not permitted.")
-            
-            json_body = {
-                "inputs" : [
-                  [{"role" : "user", "content" : prompt}]
-                ],
-                "parameters" : {
-                    "max_new_tokens" : self.max_new_tokens,
-                    "top_p" : self.top_p,
-                    "temperature" : self.temperature
-                }
-            }
-    
-            response = requests.post(model_api_gateway, json=json_body)
-    
-            return response.json()[0]["generation"]["content"]
-    
-        @property
-        def _identifying_params(self) -> Mapping[str, Any]:
-            """Get the identifying parameters."""
-            return {
-                "max_new_tokens" : self.max_new_tokens,
-                "top_p" : self.top_p,
-                "temperature" : self.temperature
-            }
-
-    # set our llm to the custom Llama2SageMaker endpoint model
-    llm = Llama2SageMaker()
-
     # Create a chain that uses the Chroma vector store
     chain = RetrievalQAWithSourcesChain.from_chain_type(
-        llm=llm,
+        ChatOpenAI(
+            model="gpt-3.5-turbo",
+            temperature=0.0
+        ),
         chain_type="stuff",
         retriever=docsearch.as_retriever(),
         return_source_documents=True,