aesth.lol
A api for Genue AI models. Self hosted.
1.5 kB
49 lines
1from transformers import pipeline
2from flask import app, request, make_response, Request
3import os
4import dotenv
5import json
6
7# LOAD ENV VARS!! VERY IMPORTANT
8dotenv.load_dotenv()
9
10_MODEL = os.environ.get("MODEL", "Geode-Onyx-2")
11pipe = pipeline("text-generation", model=f"GenueAI/{_MODEL}")
12app = app.Flask("genue-api")
13
14def _check_keys(request: Request):
15 keys = json.loads(os.environ.get("KEYS", []))
16 for i,v in enumerate(keys):
17 if request.args.get("key", "") == v:
18 return True
19 return False
20
21def _create_err(code: int):
22 if code == 401:
23 response = make_response("Invalid API key", 401)
24 response.mimetype = "text/plain"
25 return response
26 else:
27 response = make_response(f"An error occured. Code: {code}", code)
28 response.mimetype = "text/plain"
29 return response
30
31def _get_completion(prompt):
32 messages = [
33 {"role": "system", "content": f"You are an AI chatbot running on {_MODEL}"},
34 {"role": "user", "content": prompt},
35 ]
36 message_returned = pipe(messages)[0]["generated_text"][2]["content"]
37 return message_returned
38
39
40@app.route("/genue")
41def genue_model_api():
42 if _check_keys(request):
43 prompt = request.args.get("prompt", "Hello! Who are you?")
44 completion = _get_completion(prompt)
45 response = make_response(completion, 200)
46 response.mimetype = "text/plain"
47 return response
48 else:
49 return _create_err(401)