Update app.py

app.py

@@ -59,6 +59,16 @@ class DocumentRetrievalAndGeneration:
         )
         return generate_text
 
+    def generate_response_with_timeout(self, model_inputs):
+        try:
+            # Start the generation process and set a timeout
+            with multiprocessing.Pool(processes=1) as pool:
+                result = pool.apply_async(self.llm.model.generate, (model_inputs,), {"max_new_tokens": 1000, "do_sample": True})
+                generated_ids = result.get(timeout=80)  # Timeout set to 60 seconds
+            return generated_ids
+        except multiprocessing.TimeoutError:
+            raise TimeoutError("Text generation process timed out")
+    
     def query_and_generate_response(self, query):
         query_embedding = self.embeddings.encode(query, convert_to_tensor=True).cpu().numpy()
         distances, indices = self.gpu_index.search(np.array([query_embedding]), k=5)
@@ -99,7 +109,8 @@ class DocumentRetrievalAndGeneration:
         
         # Perform inference and measure time
         start_time = datetime.now()
-        generated_ids = self.llm.model.generate(model_inputs, max_new_tokens=1000, do_sample=True)
+        generated_ids = self.generate_response_with_timeout(model_inputs)
+        # generated_ids = self.llm.model.generate(model_inputs, max_new_tokens=1000, do_sample=True)
         elapsed_time = datetime.now() - start_time
 
         # Decode and return output