Update handler.py

handler.py

@@ -1,23 +1,60 @@
-from flask import Flask, request, jsonify
-from transformers import AutoTokenizer, AutoModelForCausalLM
+import json
+import os
+from typing import Dict, List, Any
 import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
 
-app = Flask(__name__)
+MAX_TOKENS=8192
+GPU_LAYERS=99 if torch.cuda.is_available() else 0
 
-# Loading
-tokenizer = AutoTokenizer.from_pretrained("MrOvkill/Phi-3-Instruct-Bloated")
-model = AutoModelForCausalLM.from_pretrained("MrOvkill/Phi-3-Instruct-Bloated")
+class EndpointHandler():
+    def __init__(self, data):
+        cfg = {
+            "repo": "MrOvkill/Phi-3-Instruct-Bloated",
+        }
+        self.model = AutoModelForCausalLM.from_pretrained(cfg['repo'])
+        self.tokenizer = AutoTokenizer.from_pretrained(cfg['repo'])
 
-@app.route('/predict', methods=['POST'])
-def predict():
-    data = request.json
-    prompt = data["prompt"]
-    kwargs = data.get('kwargs', {})
-    inputs = tokenizer(prompt, return_tensors="pt")
-    with torch.no_grad():
-        outputs = model.generate(**inputs, **kwargs)
-    response = {"response": tokenizer.decode(outputs[0], skip_special_tokens=True)}
-    return jsonify(response)
+    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
+        inputs = data.pop("inputs", "")
+        temperature = data.pop("temperature", None)
+        if not temperature:
+            temperature = data.pop("temp", 0.33)
+        if temperature > 3 or temperature < 0:
+            return json.dumps({
+                "status": "error",
+                "reason": "invalid temperature ( 0.01 - 1.00 )"
+            })
+        top_p = data.pop("top-p", 0.85)
+        if top_p > 3 or top_p < 0:
+            return json.dumps({
+                "status": "error",
+                "reason": "invalid top percentage ( 0.01 - 1.00 )"
+            })
+        top_k = data.pop("top-k", 42)
+        if top_k > 100 or top_k < 0:
+            return json.dumps({
+                "status": "error",
+                "reason": "invalid top k ( 1 - 99 )"
+            })
+        system_prompt = data.pop("system-prompt", "You are a helpful assistant.")
+        fmat = data.pop("format", f"<|system|>\n{system_prompt} <|end|>\n<|user|>\n{inputs} <|end|>\n<|assistant|>")
+        try:
+            fmat = fmat.format(system_prompt = system_prompt, prompt = inputs)
+        except Exception as e:
+            return json.dumps({
+                "status": "error",
+                "reason": "invalid format"
+            })
+        max_length = data.pop("max_length", 1024)
+        try:
+            max_length = int(max_length)
+        except Exception as e:
+            return json.dumps({
+                "status": "error",
+                "reason": "max_length was passed as something that was absolutely not a plain old int"
+            })
+        
+        res = self.model(fmat, temperature=temperature, top_p=top_p, top_k=top_k, max_tokens=max_length)
 
-if __name__ == '__main__':
-    app.run()
+        return res