MrOvkill commited on
Commit
83a5aeb
1 Parent(s): 0ffc2ad

Create handler.py

Browse files
Files changed (1) hide show
  1. handler.py +23 -0
handler.py ADDED
@@ -0,0 +1,23 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from flask import Flask, request, jsonify
2
+ from transformers import AutoTokenizer, AutoModelForCausalLM
3
+ import torch
4
+
5
+ app = Flask(__name__)
6
+
7
+ # Loading
8
+ tokenizer = AutoTokenizer.from_pretrained("MrOvkill/Phi-3-Instruct-Bloated")
9
+ model = AutoModelForCausalLM.from_pretrained("MrOvkill/Phi-3-Instruct-Bloated")
10
+
11
+ @app.route('/predict', methods=['POST'])
12
+ def predict():
13
+ data = request.json
14
+ prompt = data["prompt"]
15
+ kwargs = data.get('kwargs', {})
16
+ inputs = tokenizer(prompt, return_tensors="pt")
17
+ with torch.no_grad():
18
+ outputs = model.generate(**inputs, **kwargs)
19
+ response = {"response": tokenizer.decode(outputs[0], skip_special_tokens=True)}
20
+ return jsonify(response)
21
+
22
+ if __name__ == '__main__':
23
+ app.run()