Hivra commited on
Commit
c768d15
·
verified ·
1 Parent(s): eb281f4

Create main.py

Browse files
Files changed (1) hide show
  1. main.py +163 -0
main.py ADDED
@@ -0,0 +1,163 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gevent.pywsgi
2
+ from gevent import monkey;monkey.patch_all()
3
+ from flask import Flask, request, Response, jsonify
4
+ import argparse
5
+ import requests
6
+ import random
7
+ import string
8
+ import time
9
+ import json
10
+ import os
11
+
12
+ app = Flask(__name__)
13
+ app.json.sort_keys = False
14
+
15
+ parser = argparse.ArgumentParser(description="An example of minimax demo with a similar API to OAI.")
16
+ parser.add_argument("--host", type=str, help="Set the ip address.(default: 0.0.0.0)", default='0.0.0.0')
17
+ parser.add_argument("--port", type=int, help="Set the port.(default: 7860)", default=7860)
18
+ args = parser.parse_args()
19
+
20
+ base_url = os.getenv('MODEL_BASE_URL')
21
+
22
+ @app.route('/api/v1/models', methods=["GET", "POST"])
23
+ @app.route('/v1/models', methods=["GET", "POST"])
24
+ def model_list():
25
+ time_now = int(time.time())
26
+ model_list = {
27
+ "object": "list",
28
+ "data": [
29
+ {
30
+ "id": "MiniMax-Text-01",
31
+ "object": "model",
32
+ "created": time_now,
33
+ "owned_by": "minimax"
34
+ },
35
+ {
36
+ "id": "gpt-3.5-turbo",
37
+ "object": "model",
38
+ "created": time_now,
39
+ "owned_by": "openai"
40
+ }
41
+ ]
42
+ }
43
+ return jsonify(model_list)
44
+
45
+ @app.route("/", methods=["GET"])
46
+ def index():
47
+ return Response(f'Minimax OpenAI Compatible API<br><br>'+
48
+ f'Set "{os.getenv("SPACE_URL")}/api" as proxy (or API Domain) in your Chatbot.<br><br>'+
49
+ f'The complete API is: {os.getenv("SPACE_URL")}/api/v1/chat/completions')
50
+
51
+ @app.route("/api/v1/chat/completions", methods=["POST", "OPTIONS"])
52
+ @app.route("/v1/chat/completions", methods=["POST", "OPTIONS"])
53
+ def chat_completions():
54
+
55
+ if request.method == "OPTIONS":
56
+ return Response(
57
+ headers={
58
+ "Access-Control-Allow-Origin": "*",
59
+ "Access-Control-Allow-Headers": "*",
60
+ }
61
+ )
62
+
63
+ data = request.get_json()
64
+
65
+ # reorganize data
66
+ system = None
67
+ chat_history = []
68
+ prompt = ""
69
+
70
+ if "messages" in data:
71
+ messages = data["messages"]
72
+ message_size = len(messages)
73
+
74
+ prompt = messages[-1].get("content")
75
+
76
+ # gen a random char(11) hash
77
+ chars = string.ascii_lowercase + string.digits
78
+ session_hash = "".join(random.choice(chars) for _ in range(10))
79
+
80
+ single_prompt_data = {
81
+ 'data': [
82
+ prompt,
83
+ None,
84
+ ],
85
+ 'event_data': None,
86
+ 'fn_index': 0,
87
+ 'trigger_id': 7,
88
+ 'session_hash': session_hash,
89
+ }
90
+ response = requests.post(f'{base_url}/gradio_api/run/predict', json=single_prompt_data)
91
+
92
+ context_data = {
93
+ 'data': [
94
+ None,
95
+ messages,
96
+ 1000000,
97
+ data['temperature'],
98
+ data['top_p'],
99
+ ],
100
+ 'event_data': None,
101
+ 'fn_index': 2,
102
+ 'trigger_id': 7,
103
+ 'session_hash': session_hash,
104
+ }
105
+ response = requests.post(f'{base_url}/gradio_api/queue/join', json=context_data)
106
+
107
+ def generate():
108
+
109
+ url = f"{base_url}/gradio_api/queue/data?session_hash={session_hash}"
110
+ data = requests.get(url, stream=True)
111
+
112
+ time_now = int(time.time())
113
+
114
+ for line in data.iter_lines():
115
+ if line:
116
+ decoded_line = line.decode("utf-8")
117
+ json_line = json.loads(decoded_line[6:])
118
+ if json_line["msg"] == "process_starts":
119
+ res_data = gen_res_data({}, time_now=time_now, start=True)
120
+ yield f"data: {json.dumps(res_data)}\n\n"
121
+ elif json_line["msg"] == "process_generating":
122
+ res_data = gen_res_data(json_line, time_now=time_now)
123
+ yield f"data: {json.dumps(res_data)}\n\n"
124
+ elif json_line["msg"] == "process_completed":
125
+ yield "data: [DONE]"
126
+
127
+ return Response(
128
+ generate(),
129
+ mimetype="text/event-stream",
130
+ headers={
131
+ "Access-Control-Allow-Origin": "*",
132
+ "Access-Control-Allow-Headers": "*",
133
+ },
134
+ )
135
+
136
+
137
+ def gen_res_data(data, time_now=0, start=False):
138
+ res_data = {
139
+ "id": "chatcmpl",
140
+ "object": "chat.completion.chunk",
141
+ "created": time_now,
142
+ "model": "MiniMax-Text-01",
143
+ "choices": [{"index": 0, "finish_reason": None}],
144
+ }
145
+
146
+ if start:
147
+ res_data["choices"][0]["delta"] = {"role": "assistant", "content": ""}
148
+ else:
149
+ chat_pair = data["output"]["data"][0]
150
+
151
+ if chat_pair == []:
152
+ res_data["choices"][0]["finish_reason"] = "stop"
153
+ else:
154
+ try:
155
+ res_data["choices"][0]["delta"] = {"content": chat_pair[-1][-1]}
156
+ except:
157
+ res_data["choices"][0]["delta"] = {"content": chat_pair[-1]["content"]}
158
+ return res_data
159
+
160
+
161
+ if __name__ == "__main__":
162
+ #app.run(host=args.host, port=args.port, debug=True)
163
+ gevent.pywsgi.WSGIServer((args.host, args.port), app).serve_forever()