jdgalvan commited on
Commit
e06a121
·
verified ·
1 Parent(s): f98e8e2

Delete handler.py

Browse files
Files changed (1) hide show
  1. handler.py +0 -32
handler.py DELETED
@@ -1,32 +0,0 @@
1
- from typing import Dict, List, Any
2
- from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
3
-
4
- class EndpointHandler():
5
- def __init__(self, path=""):
6
- # Load model directly
7
- model = AutoModelForCausalLM.from_pretrained(
8
- "jdgalvan/Phi-3-mini-4k-instruct",
9
- device_map="cuda",
10
- torch_dtype="auto",
11
- trust_remote_code=True,
12
- )
13
-
14
- tokenizer = AutoTokenizer.from_pretrained("jdgalvan/Phi-3-mini-4k-instruct")
15
-
16
- self.pipe = pipeline(
17
- "text-generation",
18
- model=model,
19
- tokenizer=tokenizer,
20
- )
21
-
22
- def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
23
- inputs = data.pop("inputs", data)
24
- parameters = data.pop("parameters", None)
25
-
26
- # pass inputs with all kwargs in data
27
- if parameters is not None:
28
- prediction = self.pipe(inputs, **parameters)
29
- else:
30
- prediction = self.pipe(inputs)
31
-
32
- return prediction