|
|
import json |
|
|
import argparse |
|
|
import torch |
|
|
import os |
|
|
import random |
|
|
import numpy as np |
|
|
import requests |
|
|
import logging |
|
|
import math |
|
|
import copy |
|
|
import string |
|
|
|
|
|
from tqdm import tqdm |
|
|
from time import time |
|
|
from flask import Flask, request, jsonify |
|
|
from flask_cors import CORS |
|
|
from tornado.wsgi import WSGIContainer |
|
|
from tornado.httpserver import HTTPServer |
|
|
from tornado.ioloop import IOLoop |
|
|
|
|
|
from simcse import SimCSE |
|
|
|
|
|
logging.basicConfig(format='%(asctime)s - %(levelname)s - %(name)s - %(message)s', datefmt='%m/%d/%Y %H:%M:%S', |
|
|
level=logging.INFO) |
|
|
logger = logging.getLogger(__name__) |
|
|
|
|
|
def run_simcse_demo(port, args): |
|
|
app = Flask(__name__, static_folder='./static') |
|
|
app.config['JSONIFY_PRETTYPRINT_REGULAR'] = False |
|
|
CORS(app) |
|
|
|
|
|
sentence_path = os.path.join(args.sentences_dir, args.example_sentences) |
|
|
query_path = os.path.join(args.sentences_dir, args.example_query) |
|
|
embedder = SimCSE(args.model_name_or_path) |
|
|
embedder.build_index(sentence_path) |
|
|
@app.route('/') |
|
|
def index(): |
|
|
return app.send_static_file('index.html') |
|
|
|
|
|
@app.route('/api', methods=['GET']) |
|
|
def api(): |
|
|
query = request.args['query'] |
|
|
top_k = int(request.args['topk']) |
|
|
threshold = float(request.args['threshold']) |
|
|
start = time() |
|
|
results = embedder.search(query, top_k=top_k, threshold=threshold) |
|
|
ret = [] |
|
|
out = {} |
|
|
for sentence, score in results: |
|
|
ret.append({"sentence": sentence, "score": score}) |
|
|
span = time() - start |
|
|
out['ret'] = ret |
|
|
out['time'] = "{:.4f}".format(span) |
|
|
return jsonify(out) |
|
|
|
|
|
@app.route('/files/<path:path>') |
|
|
def static_files(path): |
|
|
return app.send_static_file('files/' + path) |
|
|
|
|
|
@app.route('/get_examples', methods=['GET']) |
|
|
def get_examples(): |
|
|
with open(query_path, 'r') as fp: |
|
|
examples = [line.strip() for line in fp.readlines()] |
|
|
return jsonify(examples) |
|
|
|
|
|
addr = args.ip + ":" + args.port |
|
|
logger.info(f'Starting Index server at {addr}') |
|
|
http_server = HTTPServer(WSGIContainer(app)) |
|
|
http_server.listen(port) |
|
|
IOLoop.instance().start() |
|
|
|
|
|
if __name__=="__main__": |
|
|
parser = argparse.ArgumentParser() |
|
|
parser.add_argument('--model_name_or_path', default=None, type=str) |
|
|
parser.add_argument('--device', default='cpu', type=str) |
|
|
parser.add_argument('--sentences_dir', default=None, type=str) |
|
|
parser.add_argument('--example_query', default=None, type=str) |
|
|
parser.add_argument('--example_sentences', default=None, type=str) |
|
|
parser.add_argument('--port', default='8888', type=str) |
|
|
parser.add_argument('--ip', default='http://127.0.0.1') |
|
|
parser.add_argument('--load_light', default=False, action='store_true') |
|
|
args = parser.parse_args() |
|
|
|
|
|
run_simcse_demo(args.port, args) |
