Update handler.py
Browse files- handler.py +0 -5
handler.py
CHANGED
@@ -2,10 +2,6 @@ import torch
|
|
2 |
from typing import Dict, List, Any
|
3 |
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
|
4 |
|
5 |
-
# check for GPU
|
6 |
-
device = 0 if torch.cuda.is_available() else -1
|
7 |
-
|
8 |
-
|
9 |
format_input = (
|
10 |
"Below is an instruction that describes a task. "
|
11 |
"Write a response that appropriately completes the request.\n\n"
|
@@ -27,7 +23,6 @@ class EndpointHandler:
|
|
27 |
"text-generation",
|
28 |
model=model,
|
29 |
tokenizer=tokenizer,
|
30 |
-
device=device,
|
31 |
max_length=256,
|
32 |
)
|
33 |
|
|
|
2 |
from typing import Dict, List, Any
|
3 |
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
|
4 |
|
|
|
|
|
|
|
|
|
5 |
format_input = (
|
6 |
"Below is an instruction that describes a task. "
|
7 |
"Write a response that appropriately completes the request.\n\n"
|
|
|
23 |
"text-generation",
|
24 |
model=model,
|
25 |
tokenizer=tokenizer,
|
|
|
26 |
max_length=256,
|
27 |
)
|
28 |
|