nikita
commited on
Commit
•
b66e7ec
1
Parent(s):
e67a0a3
remove mps device in handler
Browse files- handler.py +1 -3
handler.py
CHANGED
@@ -14,9 +14,7 @@ class LLMLinguaParameters(BaseModel):
|
|
14 |
class EndpointHandler:
|
15 |
def __init__(self, path=None):
|
16 |
model_path = path or "NousResearch/Llama-2-7b-hf"
|
17 |
-
|
18 |
-
# self.model = None # TODO: load model
|
19 |
-
self.llm_lingua = PromptCompressor(model_name=model_path, device_map="mps") # TODO: remove mps device
|
20 |
|
21 |
def __call__(self, data: dict[str, Any]) -> list[dict[str, Any]]:
|
22 |
"""
|
|
|
14 |
class EndpointHandler:
|
15 |
def __init__(self, path=None):
|
16 |
model_path = path or "NousResearch/Llama-2-7b-hf"
|
17 |
+
self.llm_lingua = PromptCompressor(model_name=model_path)
|
|
|
|
|
18 |
|
19 |
def __call__(self, data: dict[str, Any]) -> list[dict[str, Any]]:
|
20 |
"""
|