Upload folder using huggingface_hub

handler.py

@@ -28,15 +28,15 @@ class EndpointHandler():
         # run generation
         samples = []
         with torch.autocast(device_type=self.device_type, dtype=torch.bfloat16):
-            for k in range(data['parameters']['num_samples']):
-                for _ in range(data['parameters']['max_new_tokens']):
+            for k in range(int(data['parameters']['num_samples'])):
+                for _ in range(int(data['parameters']['max_new_tokens'])):
                     # forward the model to get the logits for the index in the sequence
                     logits, _ = self.model(idx)
                     # pluck the logits at the final step and scale by desired temperature
                     logits = logits[:, -1, :] / data['parameters']['temperature']
                     # optionally crop the logits to only the top k options
-                    if data['parameters']['top_k'] is not None:
-                        v, _ = torch.topk(logits, min(data['parameters']['top_k'], logits.size(-1)))
+                    if int(data['parameters']['top_k']) is not None:
+                        v, _ = torch.topk(logits, min(int(data['parameters']['top_k']), logits.size(-1)))
                         logits[logits < v[:, [-1]]] = -float('Inf')
                     # apply softmax to convert logits to (normalized) probabilities
                     probs = torch.nn.functional.softmax(logits, dim=-1)