Update handler.py

handler.py

@@ -48,11 +48,22 @@ class EndpointHandler:
             # Tokenize input text
             input_ids = self.tokenizer(user_text, return_tensors="pt").input_ids
 
+            # Measure inference time
+            start_time = time.time()
+            
             # Perform inference
             with torch.no_grad():
                 output_ids = self.model.generate(input_ids, max_length=100, temperature=0.3)
 
             json_output = self.tokenizer.decode(output_ids[0], skip_special_tokens=True)
+
+            end_time = time.time()
+            inference_time = end_time - start_time  # Calculate time taken
+
+            # Print inference time
+            print(f"Inference Time: {inference_time:.4f} seconds")
+
+        
             # return json.loads(json_output)
             try:
                 return json.loads(json_output)