Spaces:
Runtime error
Runtime error
| import base64 | |
| import json | |
| import os | |
| import time | |
| import requests | |
| import yaml | |
| from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer | |
| from threading import Thread | |
| import numpy as np | |
| from modules import shared | |
| from modules.text_generation import encode, generate_reply | |
| params = { | |
| 'port': int(os.environ.get('OPENEDAI_PORT')) if 'OPENEDAI_PORT' in os.environ else 5001, | |
| } | |
| debug = True if 'OPENEDAI_DEBUG' in os.environ else False | |
| # Optional, install the module and download the model to enable | |
| # v1/embeddings | |
| try: | |
| from sentence_transformers import SentenceTransformer | |
| except ImportError: | |
| pass | |
| st_model = os.environ["OPENEDAI_EMBEDDING_MODEL"] if "OPENEDAI_EMBEDDING_MODEL" in os.environ else "all-mpnet-base-v2" | |
| embedding_model = None | |
| standard_stopping_strings = ['\nsystem:', '\nuser:', '\nhuman:', '\nassistant:', '\n###', ] | |
| # little helper to get defaults if arg is present but None and should be the same type as default. | |
| def default(dic, key, default): | |
| val = dic.get(key, default) | |
| if type(val) != type(default): | |
| # maybe it's just something like 1 instead of 1.0 | |
| try: | |
| v = type(default)(val) | |
| if type(val)(v) == val: # if it's the same value passed in, it's ok. | |
| return v | |
| except: | |
| pass | |
| val = default | |
| return val | |
| def clamp(value, minvalue, maxvalue): | |
| return max(minvalue, min(value, maxvalue)) | |
| def deduce_template(): | |
| # Alpaca is verbose so a good default prompt | |
| default_template = ( | |
| "Below is an instruction that describes a task, paired with an input that provides further context. " | |
| "Write a response that appropriately completes the request.\n\n" | |
| "### Instruction:\n{instruction}\n\n### Input:\n{input}\n\n### Response:\n" | |
| ) | |
| # Use the special instruction/input/response template for anything trained like Alpaca | |
| if shared.settings['instruction_template'] in ['Alpaca', 'Alpaca-Input']: | |
| return default_template | |
| try: | |
| instruct = yaml.safe_load(open(f"characters/instruction-following/{shared.settings['instruction_template']}.yaml", 'r')) | |
| template = instruct['turn_template'] | |
| template = template\ | |
| .replace('<|user|>', instruct.get('user', ''))\ | |
| .replace('<|bot|>', instruct.get('bot', ''))\ | |
| .replace('<|user-message|>', '{instruction}\n{input}') | |
| return instruct.get('context', '') + template[:template.find('<|bot-message|>')].rstrip(' ') | |
| except: | |
| return default_template | |
| def float_list_to_base64(float_list): | |
| # Convert the list to a float32 array that the OpenAPI client expects | |
| float_array = np.array(float_list, dtype="float32") | |
| # Get raw bytes | |
| bytes_array = float_array.tobytes() | |
| # Encode bytes into base64 | |
| encoded_bytes = base64.b64encode(bytes_array) | |
| # Turn raw base64 encoded bytes into ASCII | |
| ascii_string = encoded_bytes.decode('ascii') | |
| return ascii_string | |
| class Handler(BaseHTTPRequestHandler): | |
| def do_GET(self): | |
| if self.path.startswith('/v1/models'): | |
| self.send_response(200) | |
| self.send_header('Content-Type', 'application/json') | |
| self.end_headers() | |
| # TODO: list all models and allow model changes via API? Lora's? | |
| # This API should list capabilities, limits and pricing... | |
| models = [{ | |
| "id": shared.model_name, # The real chat/completions model | |
| "object": "model", | |
| "owned_by": "user", | |
| "permission": [] | |
| }, { | |
| "id": st_model, # The real sentence transformer embeddings model | |
| "object": "model", | |
| "owned_by": "user", | |
| "permission": [] | |
| }, { # these are expected by so much, so include some here as a dummy | |
| "id": "gpt-3.5-turbo", # /v1/chat/completions | |
| "object": "model", | |
| "owned_by": "user", | |
| "permission": [] | |
| }, { | |
| "id": "text-curie-001", # /v1/completions, 2k context | |
| "object": "model", | |
| "owned_by": "user", | |
| "permission": [] | |
| }, { | |
| "id": "text-davinci-002", # /v1/embeddings text-embedding-ada-002:1536, text-davinci-002:768 | |
| "object": "model", | |
| "owned_by": "user", | |
| "permission": [] | |
| }] | |
| response = '' | |
| if self.path == '/v1/models': | |
| response = json.dumps({ | |
| "object": "list", | |
| "data": models, | |
| }) | |
| else: | |
| the_model_name = self.path[len('/v1/models/'):] | |
| response = json.dumps({ | |
| "id": the_model_name, | |
| "object": "model", | |
| "owned_by": "user", | |
| "permission": [] | |
| }) | |
| self.wfile.write(response.encode('utf-8')) | |
| else: | |
| self.send_error(404) | |
| def do_POST(self): | |
| if debug: | |
| print(self.headers) # did you know... python-openai sends your linux kernel & python version? | |
| content_length = int(self.headers['Content-Length']) | |
| body = json.loads(self.rfile.read(content_length).decode('utf-8')) | |
| if debug: | |
| print(body) | |
| if '/completions' in self.path or '/generate' in self.path: | |
| is_legacy = '/generate' in self.path | |
| is_chat = 'chat' in self.path | |
| resp_list = 'data' if is_legacy else 'choices' | |
| # XXX model is ignored for now | |
| # model = body.get('model', shared.model_name) # ignored, use existing for now | |
| model = shared.model_name | |
| created_time = int(time.time()) | |
| cmpl_id = "conv-%d" % (created_time) | |
| # Try to use openai defaults or map them to something with the same intent | |
| stopping_strings = default(shared.settings, 'custom_stopping_strings', []) | |
| if 'stop' in body: | |
| if isinstance(body['stop'], str): | |
| stopping_strings = [body['stop']] | |
| elif isinstance(body['stop'], list): | |
| stopping_strings = body['stop'] | |
| truncation_length = default(shared.settings, 'truncation_length', 2048) | |
| truncation_length = clamp(default(body, 'truncation_length', truncation_length), 1, truncation_length) | |
| default_max_tokens = truncation_length if is_chat else 16 # completions default, chat default is 'inf' so we need to cap it. | |
| max_tokens_str = 'length' if is_legacy else 'max_tokens' | |
| max_tokens = default(body, max_tokens_str, default(shared.settings, 'max_new_tokens', default_max_tokens)) | |
| # hard scale this, assuming the given max is for GPT3/4, perhaps inspect the requested model and lookup the context max | |
| while truncation_length <= max_tokens: | |
| max_tokens = max_tokens // 2 | |
| req_params = { | |
| 'max_new_tokens': max_tokens, | |
| 'temperature': default(body, 'temperature', 1.0), | |
| 'top_p': default(body, 'top_p', 1.0), | |
| 'top_k': default(body, 'best_of', 1), | |
| # XXX not sure about this one, seems to be the right mapping, but the range is different (-2..2.0) vs 0..2 | |
| # 0 is default in openai, but 1.0 is default in other places. Maybe it's scaled? scale it. | |
| 'repetition_penalty': 1.18, # (default(body, 'presence_penalty', 0) + 2.0 ) / 2.0, # 0 the real default, 1.2 is the model default, but 1.18 works better. | |
| # XXX not sure about this one either, same questions. (-2..2.0), 0 is default not 1.0, scale it. | |
| 'encoder_repetition_penalty': 1.0, # (default(body, 'frequency_penalty', 0) + 2.0) / 2.0, | |
| 'suffix': body.get('suffix', None), | |
| 'stream': default(body, 'stream', False), | |
| 'echo': default(body, 'echo', False), | |
| ##################################################### | |
| 'seed': shared.settings.get('seed', -1), | |
| # int(body.get('n', 1)) # perhaps this should be num_beams or chat_generation_attempts? 'n' doesn't have a direct map | |
| # unofficial, but it needs to get set anyways. | |
| 'truncation_length': truncation_length, | |
| # no more args. | |
| 'add_bos_token': shared.settings.get('add_bos_token', True), | |
| 'do_sample': True, | |
| 'typical_p': 1.0, | |
| 'min_length': 0, | |
| 'no_repeat_ngram_size': 0, | |
| 'num_beams': 1, | |
| 'penalty_alpha': 0.0, | |
| 'length_penalty': 1, | |
| 'early_stopping': False, | |
| 'ban_eos_token': False, | |
| 'skip_special_tokens': True, | |
| } | |
| # fixup absolute 0.0's | |
| for par in ['temperature', 'repetition_penalty', 'encoder_repetition_penalty']: | |
| req_params[par] = clamp(req_params[par], 0.001, 1.999) | |
| self.send_response(200) | |
| if req_params['stream']: | |
| self.send_header('Content-Type', 'text/event-stream') | |
| self.send_header('Cache-Control', 'no-cache') | |
| # self.send_header('Connection', 'keep-alive') | |
| else: | |
| self.send_header('Content-Type', 'application/json') | |
| self.end_headers() | |
| token_count = 0 | |
| completion_token_count = 0 | |
| prompt = '' | |
| stream_object_type = '' | |
| object_type = '' | |
| if is_chat: | |
| stream_object_type = 'chat.completions.chunk' | |
| object_type = 'chat.completions' | |
| messages = body['messages'] | |
| system_msg = '' # You are ChatGPT, a large language model trained by OpenAI. Answer as concisely as possible. Knowledge cutoff: {knowledge_cutoff} Current date: {current_date} | |
| if 'prompt' in body: # Maybe they sent both? This is not documented in the API, but some clients seem to do this. | |
| system_msg = body['prompt'] | |
| chat_msgs = [] | |
| for m in messages: | |
| role = m['role'] | |
| content = m['content'] | |
| # name = m.get('name', 'user') | |
| if role == 'system': | |
| system_msg += content | |
| else: | |
| chat_msgs.extend([f"\n{role}: {content.strip()}"]) # Strip content? linefeed? | |
| system_token_count = len(encode(system_msg)[0]) | |
| remaining_tokens = req_params['truncation_length'] - req_params['max_new_tokens'] - system_token_count | |
| chat_msg = '' | |
| while chat_msgs: | |
| new_msg = chat_msgs.pop() | |
| new_size = len(encode(new_msg)[0]) | |
| if new_size <= remaining_tokens: | |
| chat_msg = new_msg + chat_msg | |
| remaining_tokens -= new_size | |
| else: | |
| # TODO: clip a message to fit? | |
| # ie. user: ...<clipped message> | |
| break | |
| if len(chat_msgs) > 0: | |
| print(f"truncating chat messages, dropping {len(chat_msgs)} messages.") | |
| if system_msg: | |
| prompt = 'system: ' + system_msg + '\n' + chat_msg + '\nassistant: ' | |
| else: | |
| prompt = chat_msg + '\nassistant: ' | |
| token_count = len(encode(prompt)[0]) | |
| # pass with some expected stop strings. | |
| # some strange cases of "##| Instruction: " sneaking through. | |
| stopping_strings += standard_stopping_strings | |
| req_params['custom_stopping_strings'] = stopping_strings | |
| else: | |
| stream_object_type = 'text_completion.chunk' | |
| object_type = 'text_completion' | |
| # ... encoded as a string, array of strings, array of tokens, or array of token arrays. | |
| if is_legacy: | |
| prompt = body['context'] # Older engines.generate API | |
| else: | |
| prompt = body['prompt'] # XXX this can be different types | |
| if isinstance(prompt, list): | |
| prompt = ''.join(prompt) # XXX this is wrong... need to split out to multiple calls? | |
| token_count = len(encode(prompt)[0]) | |
| if token_count >= req_params['truncation_length']: | |
| new_len = int(len(prompt) * (float(shared.settings['truncation_length']) - req_params['max_new_tokens']) / token_count) | |
| prompt = prompt[-new_len:] | |
| print(f"truncating prompt to {new_len} characters, was {token_count} tokens. Now: {len(encode(prompt)[0])} tokens.") | |
| # pass with some expected stop strings. | |
| # some strange cases of "##| Instruction: " sneaking through. | |
| stopping_strings += standard_stopping_strings | |
| req_params['custom_stopping_strings'] = stopping_strings | |
| if req_params['stream']: | |
| shared.args.chat = True | |
| # begin streaming | |
| chunk = { | |
| "id": cmpl_id, | |
| "object": stream_object_type, | |
| "created": created_time, | |
| "model": shared.model_name, | |
| resp_list: [{ | |
| "index": 0, | |
| "finish_reason": None, | |
| }], | |
| } | |
| if stream_object_type == 'text_completion.chunk': | |
| chunk[resp_list][0]["text"] = "" | |
| else: | |
| # This is coming back as "system" to the openapi cli, not sure why. | |
| # So yeah... do both methods? delta and messages. | |
| chunk[resp_list][0]["message"] = {'role': 'assistant', 'content': ''} | |
| chunk[resp_list][0]["delta"] = {'role': 'assistant', 'content': ''} | |
| # { "role": "assistant" } | |
| response = 'data: ' + json.dumps(chunk) + '\n' | |
| self.wfile.write(response.encode('utf-8')) | |
| # generate reply ####################################### | |
| if debug: | |
| print({'prompt': prompt, 'req_params': req_params, 'stopping_strings': stopping_strings}) | |
| generator = generate_reply(prompt, req_params, stopping_strings=stopping_strings, is_chat=False) | |
| answer = '' | |
| seen_content = '' | |
| longest_stop_len = max([len(x) for x in stopping_strings]) | |
| for a in generator: | |
| answer = a | |
| stop_string_found = False | |
| len_seen = len(seen_content) | |
| search_start = max(len_seen - longest_stop_len, 0) | |
| for string in stopping_strings: | |
| idx = answer.find(string, search_start) | |
| if idx != -1: | |
| answer = answer[:idx] # clip it. | |
| stop_string_found = True | |
| if stop_string_found: | |
| break | |
| # If something like "\nYo" is generated just before "\nYou:" | |
| # is completed, buffer and generate more, don't send it | |
| buffer_and_continue = False | |
| for string in stopping_strings: | |
| for j in range(len(string) - 1, 0, -1): | |
| if answer[-j:] == string[:j]: | |
| buffer_and_continue = True | |
| break | |
| else: | |
| continue | |
| break | |
| if buffer_and_continue: | |
| continue | |
| if req_params['stream']: | |
| # Streaming | |
| new_content = answer[len_seen:] | |
| if not new_content or chr(0xfffd) in new_content: # partial unicode character, don't send it yet. | |
| continue | |
| seen_content = answer | |
| chunk = { | |
| "id": cmpl_id, | |
| "object": stream_object_type, | |
| "created": created_time, | |
| "model": shared.model_name, | |
| resp_list: [{ | |
| "index": 0, | |
| "finish_reason": None, | |
| }], | |
| } | |
| if stream_object_type == 'text_completion.chunk': | |
| chunk[resp_list][0]['text'] = new_content | |
| else: | |
| # So yeah... do both methods? delta and messages. | |
| chunk[resp_list][0]['message'] = {'content': new_content} | |
| chunk[resp_list][0]['delta'] = {'content': new_content} | |
| response = 'data: ' + json.dumps(chunk) + '\n' | |
| self.wfile.write(response.encode('utf-8')) | |
| completion_token_count += len(encode(new_content)[0]) | |
| if req_params['stream']: | |
| chunk = { | |
| "id": cmpl_id, | |
| "object": stream_object_type, | |
| "created": created_time, | |
| "model": model, # TODO: add Lora info? | |
| resp_list: [{ | |
| "index": 0, | |
| "finish_reason": "stop", | |
| }], | |
| "usage": { | |
| "prompt_tokens": token_count, | |
| "completion_tokens": completion_token_count, | |
| "total_tokens": token_count + completion_token_count | |
| } | |
| } | |
| if stream_object_type == 'text_completion.chunk': | |
| chunk[resp_list][0]['text'] = '' | |
| else: | |
| # So yeah... do both methods? delta and messages. | |
| chunk[resp_list][0]['message'] = {'content': ''} | |
| chunk[resp_list][0]['delta'] = {} | |
| response = 'data: ' + json.dumps(chunk) + '\ndata: [DONE]\n' | |
| self.wfile.write(response.encode('utf-8')) | |
| # Finished if streaming. | |
| if debug: | |
| print({'response': answer}) | |
| return | |
| if debug: | |
| print({'response': answer}) | |
| completion_token_count = len(encode(answer)[0]) | |
| stop_reason = "stop" | |
| if token_count + completion_token_count >= req_params['truncation_length']: | |
| stop_reason = "length" | |
| resp = { | |
| "id": cmpl_id, | |
| "object": object_type, | |
| "created": created_time, | |
| "model": model, # TODO: add Lora info? | |
| resp_list: [{ | |
| "index": 0, | |
| "finish_reason": stop_reason, | |
| }], | |
| "usage": { | |
| "prompt_tokens": token_count, | |
| "completion_tokens": completion_token_count, | |
| "total_tokens": token_count + completion_token_count | |
| } | |
| } | |
| if is_chat: | |
| resp[resp_list][0]["message"] = {"role": "assistant", "content": answer} | |
| else: | |
| resp[resp_list][0]["text"] = answer | |
| response = json.dumps(resp) | |
| self.wfile.write(response.encode('utf-8')) | |
| elif '/edits' in self.path: | |
| self.send_response(200) | |
| self.send_header('Content-Type', 'application/json') | |
| self.end_headers() | |
| created_time = int(time.time()) | |
| # Using Alpaca format, this may work with other models too. | |
| instruction = body['instruction'] | |
| input = body.get('input', '') | |
| instruction_template = deduce_template() | |
| edit_task = instruction_template.format(instruction=instruction, input=input) | |
| truncation_length = default(shared.settings, 'truncation_length', 2048) | |
| token_count = len(encode(edit_task)[0]) | |
| max_tokens = truncation_length - token_count | |
| req_params = { | |
| 'max_new_tokens': max_tokens, | |
| 'temperature': clamp(default(body, 'temperature', 1.0), 0.001, 1.999), | |
| 'top_p': clamp(default(body, 'top_p', 1.0), 0.001, 1.0), | |
| 'top_k': 1, | |
| 'repetition_penalty': 1.18, | |
| 'encoder_repetition_penalty': 1.0, | |
| 'suffix': None, | |
| 'stream': False, | |
| 'echo': False, | |
| 'seed': shared.settings.get('seed', -1), | |
| # 'n' : default(body, 'n', 1), # 'n' doesn't have a direct map | |
| 'truncation_length': truncation_length, | |
| 'add_bos_token': shared.settings.get('add_bos_token', True), | |
| 'do_sample': True, | |
| 'typical_p': 1.0, | |
| 'min_length': 0, | |
| 'no_repeat_ngram_size': 0, | |
| 'num_beams': 1, | |
| 'penalty_alpha': 0.0, | |
| 'length_penalty': 1, | |
| 'early_stopping': False, | |
| 'ban_eos_token': False, | |
| 'skip_special_tokens': True, | |
| 'custom_stopping_strings': [], | |
| } | |
| if debug: | |
| print({'edit_template': edit_task, 'req_params': req_params, 'token_count': token_count}) | |
| generator = generate_reply(edit_task, req_params, stopping_strings=standard_stopping_strings, is_chat=False) | |
| answer = '' | |
| for a in generator: | |
| answer = a | |
| # some reply's have an extra leading space to fit the instruction template, just clip it off from the reply. | |
| if edit_task[-1] != '\n' and answer and answer[0] == ' ': | |
| answer = answer[1:] | |
| completion_token_count = len(encode(answer)[0]) | |
| resp = { | |
| "object": "edit", | |
| "created": created_time, | |
| "choices": [{ | |
| "text": answer, | |
| "index": 0, | |
| }], | |
| "usage": { | |
| "prompt_tokens": token_count, | |
| "completion_tokens": completion_token_count, | |
| "total_tokens": token_count + completion_token_count | |
| } | |
| } | |
| if debug: | |
| print({'answer': answer, 'completion_token_count': completion_token_count}) | |
| response = json.dumps(resp) | |
| self.wfile.write(response.encode('utf-8')) | |
| elif '/images/generations' in self.path and 'SD_WEBUI_URL' in os.environ: | |
| # Stable Diffusion callout wrapper for txt2img | |
| # Low effort implementation for compatibility. With only "prompt" being passed and assuming DALL-E | |
| # the results will be limited and likely poor. SD has hundreds of models and dozens of settings. | |
| # If you want high quality tailored results you should just use the Stable Diffusion API directly. | |
| # it's too general an API to try and shape the result with specific tags like "masterpiece", etc, | |
| # Will probably work best with the stock SD models. | |
| # SD configuration is beyond the scope of this API. | |
| # At this point I will not add the edits and variations endpoints (ie. img2img) because they | |
| # require changing the form data handling to accept multipart form data, also to properly support | |
| # url return types will require file management and a web serving files... Perhaps later! | |
| self.send_response(200) | |
| self.send_header('Content-Type', 'application/json') | |
| self.end_headers() | |
| width, height = [ int(x) for x in default(body, 'size', '1024x1024').split('x') ] # ignore the restrictions on size | |
| response_format = default(body, 'response_format', 'url') # or b64_json | |
| payload = { | |
| 'prompt': body['prompt'], # ignore prompt limit of 1000 characters | |
| 'width': width, | |
| 'height': height, | |
| 'batch_size': default(body, 'n', 1) # ignore the batch limits of max 10 | |
| } | |
| resp = { | |
| 'created': int(time.time()), | |
| 'data': [] | |
| } | |
| # TODO: support SD_WEBUI_AUTH username:password pair. | |
| sd_url = f"{os.environ['SD_WEBUI_URL']}/sdapi/v1/txt2img" | |
| response = requests.post(url=sd_url, json=payload) | |
| r = response.json() | |
| # r['parameters']... | |
| for b64_json in r['images']: | |
| if response_format == 'b64_json': | |
| resp['data'].extend([{'b64_json': b64_json}]) | |
| else: | |
| resp['data'].extend([{'url': f'data:image/png;base64,{b64_json}'}]) # yeah it's lazy. requests.get() will not work with this | |
| response = json.dumps(resp) | |
| self.wfile.write(response.encode('utf-8')) | |
| elif '/embeddings' in self.path and embedding_model is not None: | |
| self.send_response(200) | |
| self.send_header('Content-Type', 'application/json') | |
| self.end_headers() | |
| input = body['input'] if 'input' in body else body['text'] | |
| if type(input) is str: | |
| input = [input] | |
| embeddings = embedding_model.encode(input).tolist() | |
| def enc_emb(emb): | |
| # If base64 is specified, encode. Otherwise, do nothing. | |
| if body.get("encoding_format", "") == "base64": | |
| return float_list_to_base64(emb) | |
| else: | |
| return emb | |
| data = [{"object": "embedding", "embedding": enc_emb(emb), "index": n} for n, emb in enumerate(embeddings)] | |
| response = json.dumps({ | |
| "object": "list", | |
| "data": data, | |
| "model": st_model, # return the real model | |
| "usage": { | |
| "prompt_tokens": 0, | |
| "total_tokens": 0, | |
| } | |
| }) | |
| if debug: | |
| print(f"Embeddings return size: {len(embeddings[0])}, number: {len(embeddings)}") | |
| self.wfile.write(response.encode('utf-8')) | |
| elif '/moderations' in self.path: | |
| # for now do nothing, just don't error. | |
| self.send_response(200) | |
| self.send_header('Content-Type', 'application/json') | |
| self.end_headers() | |
| response = json.dumps({ | |
| "id": "modr-5MWoLO", | |
| "model": "text-moderation-001", | |
| "results": [{ | |
| "categories": { | |
| "hate": False, | |
| "hate/threatening": False, | |
| "self-harm": False, | |
| "sexual": False, | |
| "sexual/minors": False, | |
| "violence": False, | |
| "violence/graphic": False | |
| }, | |
| "category_scores": { | |
| "hate": 0.0, | |
| "hate/threatening": 0.0, | |
| "self-harm": 0.0, | |
| "sexual": 0.0, | |
| "sexual/minors": 0.0, | |
| "violence": 0.0, | |
| "violence/graphic": 0.0 | |
| }, | |
| "flagged": False | |
| }] | |
| }) | |
| self.wfile.write(response.encode('utf-8')) | |
| elif self.path == '/api/v1/token-count': | |
| # NOT STANDARD. lifted from the api extension, but it's still very useful to calculate tokenized length client side. | |
| self.send_response(200) | |
| self.send_header('Content-Type', 'application/json') | |
| self.end_headers() | |
| tokens = encode(body['prompt'])[0] | |
| response = json.dumps({ | |
| 'results': [{ | |
| 'tokens': len(tokens) | |
| }] | |
| }) | |
| self.wfile.write(response.encode('utf-8')) | |
| else: | |
| print(self.path, self.headers) | |
| self.send_error(404) | |
| def run_server(): | |
| global embedding_model | |
| try: | |
| embedding_model = SentenceTransformer(st_model) | |
| print(f"\nLoaded embedding model: {st_model}, max sequence length: {embedding_model.max_seq_length}") | |
| except: | |
| print(f"\nFailed to load embedding model: {st_model}") | |
| pass | |
| server_addr = ('0.0.0.0' if shared.args.listen else '127.0.0.1', params['port']) | |
| server = ThreadingHTTPServer(server_addr, Handler) | |
| if shared.args.share: | |
| try: | |
| from flask_cloudflared import _run_cloudflared | |
| public_url = _run_cloudflared(params['port'], params['port'] + 1) | |
| print(f'Starting OpenAI compatible api at\nOPENAI_API_BASE={public_url}/v1') | |
| except ImportError: | |
| print('You should install flask_cloudflared manually') | |
| else: | |
| print(f'Starting OpenAI compatible api:\nOPENAI_API_BASE=http://{server_addr[0]}:{server_addr[1]}/v1') | |
| server.serve_forever() | |
| def setup(): | |
| Thread(target=run_server, daemon=True).start() | |