feat add time in prediction.py

prediction.py

@@ -1,6 +1,7 @@
 from functools import partial
 import os
 import re
+import time
 from xml.parsers.expat import model
 
 # https://discuss.huggingface.co/t/issues-with-sadtalker-zerogpu-spaces-inquiry-about-community-grant/110625/10
@@ -16,7 +17,6 @@ else:
 
             return wrapper
 
-
 from transformers import pipeline as hf_pipeline
 
 import litellm
@@ -59,7 +59,6 @@ class ModelPrediction:
 
     def _reset_pipeline(self, model_name):
         if self._model_name != model_name:
-            print("Resetting pipeline with model", model_name)
             self._model_name = model_name
             self._pipeline = None
 
@@ -81,15 +80,19 @@ class ModelPrediction:
                 "supported models are",
                 self.model_name2pred_func.keys(),
             )
+        
 
         prompt = prompt or self.base_prompt
         prompt = prompt.format(question=question, db_schema=db_schema)
-        print(prompt)
 
+        start_time = time.time()
         prediction = self.model_name2pred_func[model_name](prompt)
+        end_time = time.time()
         prediction["response_parsed"] = self._extract_answer_from_pred(
             prediction["response"]
         )
+        prediction['time'] = end_time - start_time
+
         return prediction