knowledge/utils/vectorizer.py

import logging
from typing import List
import numpy as np
import requests
from requests.adapters import HTTPAdapter
from urllib3.util.retry import Retry

logger = logging.getLogger(__name__)

class Vectorizer:
    _instance = None
    
    def __init__(self):
        self.api_url = "http://172.16.8.98:11434/api/embeddings"
        self.model_name = "bge-m3"
        self.session = requests.Session()
        retries = Retry(total=3, backoff_factor=1)
        self.session.mount('http://', HTTPAdapter(max_retries=retries))

    @staticmethod
    def get_embedding(text: str) -> List[float]:
        return Vectorizer()._call_ollama_api(text)

    def _initialize_model(self):
        logger.info("Initialized Ollama API client for bge-m3 model")

    @classmethod
    def get_instance(cls):
        if cls._instance is None:
            cls._instance = cls()
        return cls._instance

    def chunk_text(self, text: str, chunk_size: int = 500) -> List[str]:
        tokens = self.tokenizer.tokenize(text)
        return [self.tokenizer.convert_tokens_to_string(tokens[i:i+chunk_size]) 
               for i in range(0, len(tokens), chunk_size)]

    def encode(self, text: str, batch_size: int = 32) -> List[float]:
        if not text:
            return [0.0] * 1024  # 调整为1024维以匹配数据库
        return self._call_ollama_api(text)

    def batch_encode(self, texts: List[str], batch_size: int = 64) -> List[List[float]]:
        return [self._call_ollama_api(text) for text in texts]

    def _call_ollama_api(self, text: str) -> List[float]:
        try:
            response = self.session.post(
                self.api_url,
                json={
                    "model": self.model_name,
                    "prompt": text,
                    "options": {"embedding_only": True}
                },
                timeout=30
            )
            response.raise_for_status()
            embedding_array = np.array(response.json()["embedding"])
            l2_norm = np.linalg.norm(embedding_array)
            if l2_norm > 0:
                embedding_array = embedding_array / l2_norm
            return embedding_array.tolist()
        except Exception as e:
            logger.error(f"API请求失败: {str(e)}")
            raise

    if __name__ == '__main__':
        text ='''姓名：李XX  
            性别：女  
            年龄：55岁  
            住院号：NJZY20231102  
            主诉：突发胸痛伴呼吸困难2小时  
            现病史：患者于下午3时许突发胸痛，位于心前区，呈压榨性疼痛，伴呼吸困难，持续不缓解，无恶心呕吐及二便失禁。急诊测血压150/90mmHg，心率100次/分。心电图示II、III、aVF导联ST段弓背向上抬高0.3-0.5mV。发病前2周曾诉间断性胸闷，每次持续数分钟自行缓解。  
            既往史：高血压8年（间断服用降压药），无糖尿病史，无手术史。长期吸烟（20包/年），饮酒史5年（红酒约150g/日）。  
            体格检查：BP 155/92mmHg，心率102次/分。神志清楚，痛苦面容，双肺呼吸音清，未闻及干湿性啰音。心界不大，心率102次/分，律齐，心尖区可闻及2/6级收缩期杂音。  
            辅助检查：心电图示II、III、aVF导联ST段弓背向上抬高0.3-0.5mV；心肌酶谱：肌酸激酶同工酶（CK-MB）100U/L，肌钙蛋白T（cTnT）0.4ng/ml；心脏超声示左室下壁运动减弱。  
            诊断：急性下壁心肌梗死  
            治疗计划：  
            1. 抗血小板治疗：阿司匹林300mg嚼服，氯吡格雷300mg负荷剂量后75mg qd。  
            2. 抗凝治疗：低分子量肝素4000U皮下注射，每12小时1次。  
            3. 冠状动脉介入治疗（PCI）：急诊行冠状动脉造影，必要时行支架植入术。  
            4. 调脂治疗：阿托伐他汀20mg qn。  
            5. 血管扩张剂：硝酸甘油静脉泵入。'''
        embedding = get_embedding(text)
        print(f'生成的embedding向量:\n{embedding}')