jauntybrain
commited on
Commit
•
42c44d1
1
Parent(s):
fa4523a
Update handler.py
Browse files- handler.py +4 -8
handler.py
CHANGED
@@ -25,14 +25,10 @@ class EndpointHandler:
|
|
25 |
return_tensors="pt",).to("cuda")
|
26 |
|
27 |
# pass inputs with all kwargs in data
|
28 |
-
|
29 |
-
|
30 |
-
outputs = self.model.generate(**inputs, **parameters)
|
31 |
-
else:
|
32 |
-
with torch.autocast("cuda"):
|
33 |
-
outputs = self.model.generate(**inputs,)
|
34 |
|
35 |
# postprocess the prediction
|
36 |
-
prediction = outputs[0].cpu().numpy().
|
37 |
|
38 |
-
return
|
|
|
25 |
return_tensors="pt",).to("cuda")
|
26 |
|
27 |
# pass inputs with all kwargs in data
|
28 |
+
with torch.autocast("cuda"):
|
29 |
+
outputs = self.model.generate(**inputs, do_sample=False, max_new_tokens=400)
|
|
|
|
|
|
|
|
|
30 |
|
31 |
# postprocess the prediction
|
32 |
+
prediction = outputs[0].cpu().numpy().tobytes()
|
33 |
|
34 |
+
return prediction
|