|
from transformers import AutoModelForCausalLM, AutoTokenizer |
|
from flask import Flask, request |
|
import argparse |
|
import logging |
|
|
|
|
|
class LLMInstance: |
|
|
|
def __init__(self, model_path: str, device: str = "cuda"): |
|
|
|
self.model = AutoModelForCausalLM.from_pretrained(model_path) |
|
self.tokenizer = AutoTokenizer.from_pretrained(model_path) |
|
self.model.to(device) |
|
self.device = device |
|
|
|
def query(self, message): |
|
try: |
|
messages = [ |
|
{"role": "user", "content": message}, |
|
] |
|
encodeds = self.tokenizer.apply_chat_template(messages, return_tensors="pt") |
|
model_inputs = encodeds.to(self.device) |
|
|
|
generated_ids = self.model.generate(model_inputs, max_new_tokens=1000, do_sample=True) |
|
decoded = self.tokenizer.batch_decode(generated_ids) |
|
|
|
|
|
output = decoded[0].split("[/INST]")[1].split("</s>")[0] |
|
return { |
|
'code': 0, |
|
'ret': True, |
|
'error_msg': None, |
|
'output': output |
|
} |
|
except Exception as e: |
|
return { |
|
'code': 1, |
|
'ret': False, |
|
'error_msg': str(e), |
|
'output': None |
|
} |
|
|
|
|
|
def create_app(core): |
|
app = Flask(__name__) |
|
|
|
@app.route('/ask_llm_for_answer', methods=['POST']) |
|
def ask_llm_for_answer(): |
|
user_text = request.json['user_text'] |
|
return core.query(user_text) |
|
|
|
return app |
|
|
|
|
|
if __name__ == "__main__": |
|
parser = argparse.ArgumentParser() |
|
parser.add_argument('-m', '--model_path', required=True, default='Mistral-7B-Instruct-v0.1', help='the model path of reward model') |
|
parser.add_argument('--ip', default='0.0.0.0') |
|
parser.add_argument('-p', '--port', default=8001) |
|
parser.add_argument('--debug', action='store_true') |
|
args = parser.parse_args() |
|
|
|
if args.debug: |
|
logging.getLogger().setLevel(logging.DEBUG) |
|
else: |
|
logging.getLogger().setLevel(logging.INFO) |
|
logging.getLogger().addHandler(logging.StreamHandler()) |
|
logging.getLogger().handlers[0].setFormatter(logging.Formatter("%(message)s")) |
|
|
|
core = LLMInstance(args.model_path) |
|
app = create_app(core) |
|
app.run(host=args.ip, port=args.port) |
|
|
