correct the output format

handler.py

@@ -1,4 +1,4 @@
-from typing import Dict, Any
+from typing import Dict, Any, List
 from transformers import AutoModelForCausalLM, AutoTokenizer, TextGenerationPipeline
 import torch
 
@@ -6,30 +6,32 @@ class EndpointHandler:
     def __init__(self, path=""):
         # Load model and tokenizer
         self.tokenizer = AutoTokenizer.from_pretrained(path)
-        self.model = AutoModelForCausalLM.from_pretrained(path)
+        self.model = AutoModelForCausalLM.from_pretrained(path, device_map="auto", torch_dtype=torch.float16)
         self.pipeline = TextGenerationPipeline(
             model=self.model,
             tokenizer=self.tokenizer,
             device=0 if torch.cuda.is_available() else -1
         )
 
-    def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
+    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, str]]:
         prompt_input = data.get("inputs", "")
         vibe = data.get("vibe", "Open to All Paths")  # Default fallback
 
-        # Construct Camino-aware prompt with tone and identity
+        # Construct Camino-aware prompt
         full_prompt = (
             f"#### Human (Vibe: {vibe}): {prompt_input.strip()}\n"
             f"#### Assistant (Vela - your Camino companion):"
         )
 
-        # Default generation params, override if provided
+        # Default generation params
         generation_args = data.get("parameters", {})
         generation_args.setdefault("max_new_tokens", 1024)
         generation_args.setdefault("temperature", 0.2)
         generation_args.setdefault("top_p", 0.95)
         generation_args.setdefault("do_sample", True)
 
-        # Run generation
+        # Generate response
         outputs = self.pipeline(full_prompt, **generation_args)
-        return {"generated_text": outputs[0]["generated_text"]}
+
+        # Return in correct format
+        return [{"generated_text": outputs[0]["generated_text"]}]