binaryaaron commited on
Commit
944bc62
1 Parent(s): 3298521
Files changed (3) hide show
  1. .gitignore +10 -0
  2. handler.py +25 -0
  3. tester.py +16 -0
.gitignore ADDED
@@ -0,0 +1,10 @@
 
 
 
 
 
 
 
 
 
 
 
1
+ # Python-generated files
2
+ __pycache__/
3
+ *.py[oc]
4
+ build/
5
+ dist/
6
+ wheels/
7
+ *.egg-info
8
+
9
+ # Virtual environments
10
+ .venv
handler.py ADDED
@@ -0,0 +1,25 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from typing import Dict, List, Any
2
+ import transformers
3
+ import torch
4
+
5
+ MAX_TOKENS=8192
6
+
7
+ class EndpointHandler():
8
+ def __init__(self, path=""):
9
+ self.pipeline = transformers.pipeline(
10
+ "text-generation",
11
+ model="humane-intelligence/gemma2-9b-cpt-sealionv3-instruct-endpoint",
12
+ model_kwargs={"torch_dtype": torch.bfloat16},
13
+ device_map="auto",
14
+ )
15
+
16
+
17
+ def __call__(self, data: Any) -> List[List[Dict[str, float]]]:
18
+ inputs = data.get("inputs", data)
19
+
20
+ outputs = self.pipeline(
21
+ inputs,
22
+ max_new_tokens=256,
23
+ )
24
+ print(outputs[0]["generated_text"][-1])
25
+ return outputs
tester.py ADDED
@@ -0,0 +1,16 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from handler import EndpointHandler
2
+
3
+ if __name__ == "__main__":
4
+ # init handler
5
+ my_handler = EndpointHandler(path=".")
6
+
7
+ # prepare sample payload
8
+ messages = [
9
+ {"role": "user", "content": "Apa sentimen dari kalimat berikut ini?\nKalimat: Buku ini sangat membosankan.\nJawaban: "},
10
+ ]
11
+
12
+ # test the handler
13
+ pred=my_handler.pipeline(messages)
14
+
15
+ # show results
16
+ print(pred)