chaima01 commited on
Commit
50acfbb
·
verified ·
1 Parent(s): 64664e6

Create handler.py

Browse files
Files changed (1) hide show
  1. handler.py +29 -0
handler.py ADDED
@@ -0,0 +1,29 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # handler.py
2
+ from transformers import pipeline
3
+
4
+ # this loads the files in the repo root
5
+ generator = pipeline(
6
+ "text2text-generation",
7
+ model=".",
8
+ tokenizer="."
9
+ )
10
+
11
+ def init():
12
+ # optional: run once at container startup
13
+ pass
14
+
15
+ def inference(payload: dict):
16
+ """
17
+ payload is the body of the POST, e.g.
18
+ { "inputs": "VIBE: Agnostic\nQUESTION: …\nAnswer:" }
19
+ """
20
+ text = payload.get("inputs", "")
21
+ out = generator(
22
+ text,
23
+ max_new_tokens=128,
24
+ do_sample=True,
25
+ top_p=0.9,
26
+ temperature=0.7
27
+ )
28
+ # return whatever JSON you want clients to see
29
+ return {"generated_text": out[0]["generated_text"]}