handler v9 (#9)

- added custom handler v9 (cd11daeba612cd8893250017ccb612aaf4528395)

handler.py

@@ -1,7 +1,9 @@
 from typing import Dict, List, Any
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
+import json
 import logging
+import time
 
 # Configure logging
 logging.basicConfig(
@@ -72,12 +74,10 @@ class EndpointHandler:
                 tokenize=False,
                 add_generation_prompt=True
             )
-            logger.info(f"Generated chat prompt: {prompt}")
+            logger.info(f"Generated chat prompt: {json.dumps(prompt)}")
             
             # Tokenize the prompt
-            logger.info("Tokenizing input")
             inputs = self.tokenizer(prompt, return_tensors="pt").to(self.model.device)
-            logger.info(f"Input shape: {inputs.input_ids.shape}")
             
             # Generate response
             logger.info("Generating response")
@@ -86,7 +86,6 @@ class EndpointHandler:
                     **inputs,
                     **gen_params
                 )
-            logger.info(f"Output shape: {output_tokens.shape}")
             
             # Decode the response
             logger.info("Decoding response")
@@ -94,11 +93,29 @@ class EndpointHandler:
             
             # Extract the assistant's response by removing the input prompt
             response = output_text#[len(prompt):].strip()
-            logger.info(f"Generated response length: {len(response)}")
-            logger.info(f"Generated response: {response}")
+            logger.info(f"Generated response: {json.dumps(response)}")
             
             #return [{"role": "assistant", "content": response}]
-            return [{"result": response, "error": None}] 
+            #return {"result": response, "error": None} 
+            return {
+                "id": "cmpl-" + str(hash(response))[:10],  # Generate a unique ID
+                "object": "chat.completion",
+                "created": int(time.time()),
+                "model": self.model.config.name_or_path,
+                "choices": [{
+                    "index": 0,
+                    "message": {
+                        "role": "assistant",
+                        "content": response
+                    },
+                    "finish_reason": "stop"
+                }],
+                "usage": {
+                    "prompt_tokens": len(inputs["input_ids"][0]),
+                    "completion_tokens": len(output_tokens[0]) - len(inputs["input_ids"][0]),
+                    "total_tokens": len(output_tokens[0])
+                }
+            }
             
         except Exception as e:
             logger.error(f"Error during generation: {str(e)}", exc_info=True)