peterjandre commited on
Commit
adf4873
·
verified ·
1 Parent(s): ed125a8

Create handler.py

Browse files
Files changed (1) hide show
  1. handler.py +22 -0
handler.py ADDED
@@ -0,0 +1,22 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from transformers import AutoModelForSeq2SeqLM, AutoTokenizer, pipeline
2
+
3
+ # Global variables to cache model and tokenizer
4
+ model = None
5
+ tokenizer = None
6
+ nlp = None
7
+
8
+ def init():
9
+ global model, tokenizer, nlp
10
+ model_name_or_path = "."
11
+ model = AutoModelForSeq2SeqLM.from_pretrained(model_name_or_path)
12
+ tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)
13
+ nlp = pipeline("text2text-generation", model=model, tokenizer=tokenizer)
14
+
15
+ def inference(payload):
16
+ inputs = payload.get("inputs", "")
17
+ if not inputs:
18
+ return {"error": "No inputs provided"}
19
+
20
+ # Run generation pipeline
21
+ outputs = nlp(inputs, max_length=256)
22
+ return outputs