Update handler.py
Browse files- handler.py +1 -1
handler.py
CHANGED
|
@@ -42,7 +42,7 @@ class EndpointHandler:
|
|
| 42 |
|
| 43 |
# preprocess
|
| 44 |
device = "cuda" if torch.cuda.is_available() else "cpu"
|
| 45 |
-
inputs =
|
| 46 |
inputs = {k: v.to('cuda') for k, v in inputs.items()}
|
| 47 |
|
| 48 |
output = self.inference_model.generate(input_ids=inputs["input_ids"],pad_token_id=self.tokenizer.pad_token_id, max_new_tokens=256, do_sample=True, temperature=0.9, top_p=0.9, repetition_penalty=1.5, early_stopping=True, length_penalty = -0.3, num_beams=5, num_return_sequences=1)
|
|
|
|
| 42 |
|
| 43 |
# preprocess
|
| 44 |
device = "cuda" if torch.cuda.is_available() else "cpu"
|
| 45 |
+
inputs = tokenizer(prompt+' \n >> <assistant>:', return_tensors="pt").to(device)
|
| 46 |
inputs = {k: v.to('cuda') for k, v in inputs.items()}
|
| 47 |
|
| 48 |
output = self.inference_model.generate(input_ids=inputs["input_ids"],pad_token_id=self.tokenizer.pad_token_id, max_new_tokens=256, do_sample=True, temperature=0.9, top_p=0.9, repetition_penalty=1.5, early_stopping=True, length_penalty = -0.3, num_beams=5, num_return_sequences=1)
|