output serialization

handler.py

@@ -36,9 +36,14 @@ class EndpointHandler:
         with torch.no_grad():
             outputs = self.model(**inputs)
 
-        # Process outputs - this depends on your specific model and requirements
-        # For now, we'll just return the input as the output to fix the array format issue
-        return [{"input_text": input_text, "generated_text": outputs}]
+        # Process outputs - convert tensors to serializable format
+        # Extract the last hidden state and convert to list for JSON serialization
+        last_hidden_state = outputs.last_hidden_state
+
+        # Convert to Python list (serializable) - using the mean of the embeddings as a simple approach
+        embedding = last_hidden_state.mean(dim=1).cpu().numpy().tolist()
+
+        return [{"input_text": input_text, "embedding": embedding}]
 
 
 if __name__ == "__main__":