OliveSensorAPI/rag/src/main.py

import os
import time
import jwt

from config.config import base_dir, data_dir
from data_processing import Data_process
from pipeline import EmoLLMRAG

from langchain_openai import ChatOpenAI
from loguru import logger
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch
import streamlit as st
from openxlab.model import download
'''
	1）构建完整的 RAG pipeline。输入为用户 query，输出为 answer
	2）调用 embedding 提供的接口对 query 向量化
	3）下载基于 FAISS 预构建的 vector DB ，并检索对应信息
	4）调用 rerank 接口重排序检索内容
	5）调用 prompt 接口获取 system prompt 和 prompt template
	6）拼接 prompt 并调用模型返回结果

'''
def get_glm(temprature):
    llm = ChatOpenAI(
        model_name="glm-4",
        openai_api_base="https://open.bigmodel.cn/api/paas/v4",
        openai_api_key=generate_token("api-key"),
        streaming=False,
        temperature=temprature
    )
    return llm 

def generate_token(apikey: str, exp_seconds: int=100):
    try:
        id, secret = apikey.split(".")
    except Exception as e:
        raise Exception("invalid apikey", e)
 
    payload = {
        "api_key": id,
        "exp": int(round(time.time() * 1000)) + exp_seconds * 1000,
        "timestamp": int(round(time.time() * 1000)),
    }
 
    return jwt.encode(
        payload,
        secret,
        algorithm="HS256",
        headers={"alg": "HS256", "sign_type": "SIGN"},
    )

@st.cache_resource
def load_model():
    model_dir = os.path.join(base_dir,'../model') 
    logger.info(f'Loading model from {model_dir}')
    model = (
        AutoModelForCausalLM.from_pretrained('model', trust_remote_code=True)
        .to(torch.bfloat16)
        .cuda()
    )
    tokenizer = AutoTokenizer.from_pretrained('model', trust_remote_code=True)
    return model, tokenizer

def main(query, system_prompt=''):
    logger.info(data_dir)
    if not os.path.exists(data_dir):
         os.mkdir(data_dir)   
    dp = Data_process()
    vector_db = dp.load_vector_db()
    docs, retriever = dp.retrieve(query, vector_db, k=10)
    logger.info(f'Query: {query}')
    logger.info("Retrieve results===============================")
    for i, doc in enumerate(docs):
        logger.info(doc)
    passages,scores = dp.rerank(query, docs)
    logger.info("After reranking===============================")
    for i in range(len(scores)):
        logger.info(passages[i])
        logger.info(f'score: {str(scores[i])}')

if __name__ == "__main__":
    query = "我现在处于高三阶段，感到非常迷茫和害怕。我觉得自己从出生以来就是多余的，没有必要存在于这个世界。无论是在家庭、学校、朋友还是老师面前，我都感到被否定。我非常难过，对高考充满期望但成绩却不理想"
    main(query)
    #model = get_glm(0.7)
    #rag_obj = EmoLLMRAG(model, 3)
    #res = rag_obj.main(query)
    #logger.info(res)