chaima01 commited on
Commit
8e44141
·
verified ·
1 Parent(s): 50acfbb

Create inference.py

Browse files
Files changed (1) hide show
  1. inference.py +27 -0
inference.py ADDED
@@ -0,0 +1,27 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # inference.py
2
+
3
+ from transformers import pipeline
4
+
5
+ # This will be called once, at container startup
6
+ def init():
7
+ global generator
8
+ generator = pipeline(
9
+ "text2text-generation",
10
+ model=".",
11
+ tokenizer=".",
12
+ device=0, # GPU 0
13
+ max_length=128,
14
+ do_sample=True,
15
+ top_p=0.9,
16
+ temperature=0.7
17
+ )
18
+
19
+ # This will be called for every request
20
+ def run(request: dict):
21
+ """
22
+ Expects: { "inputs": "<your-prompt>" }
23
+ Returns: { "generated_text": "..." }
24
+ """
25
+ prompt = request.get("inputs", "")
26
+ outputs = generator(prompt)
27
+ return { "generated_text": outputs[0]["generated_text"] }