Trace2333/EasyPrompt

1import os2import spacy3from accelerate import PartialState4from accelerate.utils import set_seed5from flask import Flask, request, jsonify6 7from gpt2_generation import Translator8from gpt2_generation import generate_prompt, MODEL_CLASSES9 10os.environ["http_proxy"] = "http://127.0.0.1:7890"11os.environ["https_proxy"] = "http://127.0.0.1:7890"12 13app = Flask(__name__)14 15path_for_model = "./output/gpt2_openprompt/checkpoint-4500"16 17args = {18    "model_type": "gpt2",19    "model_name_or_path": path_for_model,20    "length": 80,21    "stop_token": None,22    "temperature": 1.0,23    "length_penalty": 1.2,24    "repetition_penalty": 1.2,25    "k": 3,26    "p": 0.9,27    "prefix": "",28    "padding_text": "",29    "xlm_language": "",30    "seed": 42,31    "use_cpu": False,32    "num_return_sequences": 1,33    "fp16": False,34    "jit": False,35}36 37distributed_state = PartialState(cpu=args["use_cpu"])38 39if args["seed"] is not None:40    set_seed(args["seed"])41 42tokenizer = None43model = None44zh_en_translator = None45nlp = None46 47def load_model_and_components():48    global tokenizer, model, zh_en_translator, nlp49 50    # Initialize the model and tokenizer51    try:52        args["model_type"] = args["model_type"].lower()53        model_class, tokenizer_class = MODEL_CLASSES[args["model_type"]]54    except KeyError:55        raise KeyError("the model {} you specified is not supported. You are welcome to add it and open a PR :)")56 57    tokenizer = tokenizer_class.from_pretrained(args["model_name_or_path"], padding_side='left')58    tokenizer.pad_token = tokenizer.eos_token59    tokenizer.mask_token = tokenizer.eos_token60    model = model_class.from_pretrained(args["model_name_or_path"])61    print("Model loaded!")62 63    # translator64    zh_en_translator = Translator("Helsinki-NLP/opus-mt-zh-en")65    print("Translator loaded!")66 67    # filter68    nlp = spacy.load('en_core_web_sm')69    print("Filter loaded!")70 71    # Set the model to the right device72    model.to(distributed_state.device)73 74    if args["fp16"]:75        model.half()76 77@app.route('/chat', methods=['POST'])78def chat():79    phrase = request.json.get('phrase')80 81    if tokenizer is None or model is None or zh_en_translator is None or nlp is None:82        load_model_and_components()83 84    messages = generate_prompt(85        prompt_text=phrase,86        args=args,87        zh_en_translator=zh_en_translator,88        nlp=nlp,89        model=model,90        tokenizer=tokenizer,91        distributed_state=distributed_state,92    )93 94    return jsonify(messages)95 96if __name__ == '__main__':97    load_model_and_components()98    app.run(host='0.0.0.0', port=10008, debug=False)99