向量检索与实时数仓强强融合，StoneData助客户实现高性能、低成本的非结构化数据分析

# 模型初始化
def __init__(self):
    if not os.path.exists(MODEL_PATH):
        os.makedirs(MODEL_PATH)
    if not os.path.exists("./paraphrase-mpnet-base-v2.zip"):
        url = 'https://public.ukp.informatik.tu-darmstadt.de/reimers/sentence-transformers/v0.2/paraphrase-mpnet-base-v2.zip'
        gdown.download(url)
    with zipfile.ZipFile('paraphrase-mpnet-base-v2.zip', 'r') as zip_ref:
        zip_ref.extractall(MODEL_PATH)
    self.model = SentenceTransformer(MODEL_PATH)


def sentence_encode(self, data):
    embedding = self.model.encode(data)
    sentence_embeddings = normalize(embedding)
    return sentence_embeddings.tolist()




# 上传 example.csv ,预处理 数据
@app.post('/load')
async def load_text(file: UploadFile = File(...), table_name: str = None):
    try:
        text = await file.read()
        fname = file.filename
        dirs = "data"
        if not os.path.exists(dirs):
            os.makedirs(dirs)
        fname_path = os.path.join(os.getcwd(), os.path.join(dirs, fname))
        with open(fname_path, 'wb') as f:
            f.write(text)
    except Exception :
        return {'status': False, 'msg': 'Failed to load data.'}
    try:
        total_num = do_load(table_name, fname_path,MODEL ,STONEDATA_CLI)
        LOGGER.info(f"Successfully loaded data, total count: {total_num}")
        return "Successfully loaded data!"
    except Exception as e:
        LOGGER.error(e)
        return {'status': False, 'msg': e}, 400


def do_load(table_name, file_dir, model, STONEDATA_CLI):
    if not table_name:
        table_name = DEFAULT_TABLE
    STONEDATA_CLI.create_table()
    title_data, text_data, sentence_embeddings = extract_features(file_dir, model)
    lens = 0
    for title, text, vector in zip(title_data, text_data, sentence_embeddings):
        st = "[" + ','.join(str(x) for x in vector) + "]"
        sql = "insert into " + table_name + " (title,contents,vector) values ('" + title + "','" + text  + "','"+ st +"');"
        try:
            STONEDATA_CLI.cursor.execute(sql)
            STONEDATA_CLI.conn.commit()
            LOGGER.debug(
                f"Insert vectors to StoneData in table: {table_name} with {len(vector)} rows")
        except Exception as e:
            LOGGER.error(f"StoneData ERROR: {e} with sql: {sql}")
            sys.exit(1)
        lens = lens + 1
    return len(lens)

@app.get('/search')
async def do_search_api(query_sentence: str = None):
    try:
        items = STONEDATA_CLI.search_vectors(query_sentence,MODEL, STONEDATA_CLI)
        res=[]
        for p, d, in items:
            dicts = {'title': p, 'content':d}
            res+=[dicts]
        LOGGER.info("Successfully searched similar text!")
        return res
    except Exception as e:
        LOGGER.error(e)
        return {'status': False, 'msg': e}, 400


def search_vectors(self,query_sentence,MODEL, STONEDATA_CLI,):
    # 转化为新向量
    vectors = MODEL.sentence_encode([query_sentence])
    st =  ','.join(str(x) for x in vectors)
    LOGGER.debug(f"Successfully search in sql: {st}")
    try:
        sql = "select title, contents from (select id, l2_distance(vector, '"+ st +"') as distance, title, contents from "+ DEFAULT_TABLE +") order by distance limit 10"
        LOGGER.debug(f"Successfully search in sql: {sql}")
        STONEDATA_CLI.cursor.execute(sql)
        results = STONEDATA_CLI.cursor.fetchall()
        LOGGER.debug(f"Successfully search in StoneData: {results}")
        return results
    except Exception as e:
        LOGGER.error(f"Failed to search StoneData in StoneData: {e}")
        sys.exit(1)