Update handler.py

handler.py

@@ -1,3 +1,46 @@
+Hugging Face's logo
+Hugging Face
+Search models, datasets, users...
+Models
+Datasets
+Spaces
+Docs
+Solutions
+Pricing
+
+
+
+
+ammarnasr
+/
+CodeGen2_1B_merged 
+
+like
+0
+Text Generation
+Transformers
+PyTorch
+codegen
+custom_code
+Inference Endpoints
+Model card
+Files and versions
+Community
+Settings
+CodeGen2_1B_merged
+/
+handler.py
+ammarnasr's picture
+ammarnasr
+Update handler.py
+5f48ffc
+7 minutes ago
+raw
+history
+blame
+edit
+delete
+1.06 kB
 from typing import Any, Dict, List
 import torch
 import transformers
@@ -13,7 +56,7 @@ class EndpointHandler:
         self.model = self.model.to(self.device)
 
 
-    def __call__(self, data: Dict[str, Any]) -> List[str]:
+    def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
         prompt = data["inputs"]
         if "config" in data:
           config = data.pop("config", None)
@@ -21,7 +64,5 @@ class EndpointHandler:
           config = {'max_new_tokens':100}
         input_ids = self.tokenizer(prompt, return_tensors="pt").input_ids.to(self.device)
         generated_ids = self.model.generate(input_ids, **config)
-        return [self.tokenizer.decode(generated_ids[0], skip_special_tokens=True)]
-        
-            
-        
+        generated_text = self.tokenizer.decode(generated_ids[0], skip_special_tokens=True)
+        return [{"generated_text": generated_text}]