isage-middleware 0.1.3.1__py3-none-any.whl → 0.2.0.0__py3-none-any.whl

sage/middleware/utils/embedding/_cohere.py

@@ -1,68 +0,0 @@
-import os
-import cohere
-import asyncio
-
-import numpy as np
-
-
-async def cohere_embed(
-        texts: list[str], api_key: str, model: str = "embed-multilingual-v3.0", input_type: str = "classification",
-        embedding_types: list[str] = ["float"]
-) -> list[float]:
-    if api_key is None:
-        api_key = os.environ.get("COHERE_API_KEY")
-    # print(api_key)
-    co = cohere.AsyncClient(api_key=api_key)
-
-    response = await co.embed(
-        texts=texts,
-        model=model,
-        input_type=input_type,
-        # embedding_types=embedding_types
-    )
-    return response.embeddings
-
-import os
-import cohere
-
-def cohere_embed_sync(
-    texts: list[str],
-    api_key: str = None,
-    model: str = "embed-multilingual-v3.0",
-    input_type: str = "classification",
-    embedding_types: list[str] = ["float"]
-) -> list[list[float]]:
-    """
-    同步版本：使用 Cohere 同步客户端生成文本 embeddings。
-
-    Args:
-        texts: 文本列表
-        api_key: Cohere API Key
-        model: 模型名称
-        input_type: 输入类型，如 classification、search_document 等
-        embedding_types: 嵌入格式（默认 float）
-
-    Returns:
-        list[list[float]]: 每个文本对应的嵌入向量
-    """
-    if api_key is None:
-        api_key = os.environ.get("COHERE_API_KEY")
-    if api_key is None:
-        raise ValueError("Cohere API key must be provided.")
-
-    co = cohere.Client(api_key=api_key)
-
-    response = co.embed(
-        texts=texts,
-        model=model,
-        input_type=input_type,
-        embedding_types=embedding_types,
-    )
-    return response.embeddings
-
-
-
-
-
-
-

sage/middleware/utils/embedding/bedrock.py

@@ -1,174 +0,0 @@
-import copy
-import os
-import json
-
-import pipmaster as pm  # Pipmaster for dynamic library install
-
-# Dependencies should be installed via requirements.txt
-# aioboto3 and tenacity are required for this module
-
-try:
-    import aioboto3
-except ImportError:
-    raise ImportError(
-        "aioboto3 package is required for AWS Bedrock embedding functionality. "
-        "Please install it via: pip install aioboto3"
-    )
-
-try:
-    from tenacity import (
-        retry,
-        stop_after_attempt,
-        wait_exponential,
-        retry_if_exception_type,
-    )
-except ImportError:
-    raise ImportError(
-        "tenacity package is required for AWS Bedrock embedding functionality. "
-        "Please install it via: pip install tenacity"
-    )
-
-import numpy as np
-
-
-
-class BedrockError(Exception):
-    """Generic error for issues related to Amazon Bedrock"""
-
-
-
-async def bedrock_embed(
-    text: str,
-    model: str = "amazon.titan-embed-text-v2:0",
-    aws_access_key_id=None,
-    aws_secret_access_key=None,
-    aws_session_token=None,
-) -> list:
-    os.environ["AWS_ACCESS_KEY_ID"] = os.environ.get(
-        "AWS_ACCESS_KEY_ID", aws_access_key_id
-    )
-    os.environ["AWS_SECRET_ACCESS_KEY"] = os.environ.get(
-        "AWS_SECRET_ACCESS_KEY", aws_secret_access_key
-    )
-    os.environ["AWS_SESSION_TOKEN"] = os.environ.get(
-        "AWS_SESSION_TOKEN", aws_session_token
-    )
-
-    session = aioboto3.Session()
-    async with session.client("bedrock-sage.kernels.runtime") as bedrock_async_client:
-        model_provider = model.split(".")[0]
-
-        if model_provider == "amazon":
-            if "v2" in model:
-                body = json.dumps({
-                    "inputText": text,
-                    "embeddingTypes": ["float"],
-                })
-            elif "v1" in model:
-                body = json.dumps({"inputText": text})
-            else:
-                raise ValueError(f"Model {model} is not supported!")
-
-            response = await bedrock_async_client.invoke_model(
-                modelId=model,
-                body=body,
-                accept="application/json",
-                contentType="application/json",
-            )
-
-            response_body = await response.get("body").json()
-            return response_body["embedding"]
-
-        elif model_provider == "cohere":
-            body = json.dumps({
-                "texts": [text],
-                "input_type": "search_document",
-                "truncate": "NONE",
-            })
-
-            response = await bedrock_async_client.invoke_model(
-                model=model,
-                body=body,
-                accept="application/json",
-                contentType="application/json",
-            )
-
-            response_body = json.loads(response.get("body").read())
-            return response_body["embeddings"][0]
-
-        else:
-            raise ValueError(f"Model provider '{model_provider}' is not supported!")
-
-
-import os
-import json
-import boto3
-
-def bedrock_embed_sync(
-    text: str,
-    model: str = "amazon.titan-embed-text-v2:0",
-    aws_access_key_id=None,
-    aws_secret_access_key=None,
-    aws_session_token=None,
-) -> list[float]:
-    """
-    同步版本：使用 AWS Bedrock 生成 embedding。
-
-    Args:
-        text: 输入文本
-        model: 模型 ID，例如 "amazon.titan-embed-text-v2:0"
-        aws_access_key_id / secret / session_token: 可选 AWS 认证信息
-
-    Returns:
-        list[float]: embedding 向量
-    """
-    # 设置 AWS 环境变量（优先从参数取）
-    if aws_access_key_id:
-        os.environ["AWS_ACCESS_KEY_ID"] = aws_access_key_id
-    if aws_secret_access_key:
-        os.environ["AWS_SECRET_ACCESS_KEY"] = aws_secret_access_key
-    if aws_session_token:
-        os.environ["AWS_SESSION_TOKEN"] = aws_session_token
-
-    bedrock_client = boto3.client("bedrock-sage.kernels.runtime")
-
-    model_provider = model.split(".")[0]
-
-    if model_provider == "amazon":
-        if "v2" in model:
-            body = json.dumps({
-                "inputText": text,
-                "embeddingTypes": ["float"],
-            })
-        elif "v1" in model:
-            body = json.dumps({"inputText": text})
-        else:
-            raise ValueError(f"Model {model} is not supported!")
-
-        response = bedrock_client.invoke_model(
-            modelId=model,
-            body=body,
-            accept="application/json",
-            contentType="application/json",
-        )
-        response_body = json.loads(response["body"].read())
-        return response_body["embedding"]
-
-    elif model_provider == "cohere":
-        body = json.dumps({
-            "texts": [text],
-            "input_type": "search_document",
-            "truncate": "NONE",
-        })
-
-        response = bedrock_client.invoke_model(
-            modelId=model,
-            body=body,
-            accept="application/json",
-            contentType="application/json",
-        )
-        response_body = json.loads(response["body"].read())
-        return response_body["embeddings"][0]
-
-    else:
-        raise ValueError(f"Model provider '{model_provider}' is not supported!")

sage/middleware/utils/embedding/embedding_api.py

@@ -1,12 +0,0 @@
-from sage.middleware.utils.embedding.embedding_model import EmbeddingModel
-
-
-def apply_embedding_model(name: str = "default",**kwargs) -> EmbeddingModel:
-    """
-    usage  参见sage/api/model/operator_test.py
-    while name(method) = "hf", please set the param:model;
-    while name(method) = "openai",if you need call other APIs which are compatible with openai,set the params:base_url,api_key,model;
-    while name(method) = "jina/siliconcloud/cohere",please set the params:api_key,model;
-    Example:operator_test.py
-    """
-    return EmbeddingModel(method=name,**kwargs)

sage/middleware/utils/embedding/embedding_model.py

@@ -1,150 +0,0 @@
-import os
-import sys
-os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"
-
-# 添加项目根目录到Python路径
-sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../../..")))
-
-import time
-
-from dotenv import load_dotenv
-
-load_dotenv()
-
-from sage.middleware.utils.embedding import hf, ollama, siliconcloud, openai, bedrock, zhipu, mockembedder  # , instructor
-from sage.middleware.utils.embedding import _cohere, nvidia_openai, lollms, jina
-from transformers import AutoModel, AutoTokenizer
-
-
-class EmbeddingModel:
-    # def __init__(self, method: str = "openai", model: str = "mistral-embed",
-    #              base_url: str = None, api_key: str = None):
-    def __init__(self, method: str = "openai", **kwargs):
-        """
-        初始化 embedding table
-        :param method: 指定使用的 embedding 方法名称，例如 "openai" 或 "cohere" 或“hf"等
-        """
-        self.init_method = method
-        self.dim = None
-        if method == "default":
-            method = "hf"
-            kwargs["model"] = "sentence-transformers/all-MiniLM-L6-v2"
-            
-        if method == "mockembedder":
-            kwargs["model"] = "mockembedder"  # 确保 model 参数存在
-            if "fixed_dim" not in kwargs:
-                kwargs["fixed_dim"] = 128  # 默认维度
-                
-        self.set_dim(kwargs["model"] )
-        self.method = method
-
-        # self.kwargs = {}
-        self.kwargs = kwargs
-        if method == "hf":
-            if "model" not in kwargs:
-                raise ValueError("hf method need model")
-            model_name = kwargs["model"]
-            self.kwargs["tokenizer"] = AutoTokenizer.from_pretrained(model_name)
-            self.kwargs["embed_model"] = AutoModel.from_pretrained(model_name, trust_remote_code=True)
-            self.kwargs.pop("model")
-        elif method == "mockembedder":
-            # 初始化 mockembedder
-            self.kwargs["embed_model"] = mockembedder.MockTextEmbedder(
-                model_name="mock-model",
-                fixed_dim=kwargs.get("fixed_dim", 128)
-            )
-        self.embed_fn = self._get_embed_function(method)
-
-    def set_dim(self, model_name):
-        """
-        :param model_name:
-        :return:
-        """
-        dimension_mapping = {
-            "mistral_embed": 1024,
-            "embed-multilingual-v3.0":1024,
-            "embed-english-v3.0": 1024,
-            "embed-english-light-v3.0": 384,
-            "embed-multilingual-light-v3.0": 384,
-            "embed-english-v2.0": 4096,
-            "embed-english-light-v2.0": 1024,
-            "embed-multilingual-v2.0": 768,
-            "jina-embeddings-v3":1024,
-            "BAAI/bge-m3":1024,
-            "sentence-transformers/all-MiniLM-L6-v2":384,
-            "mockembedder": 128
-        }
-        if model_name in dimension_mapping:
-            self.dim = dimension_mapping[model_name]
-        else:
-            raise ValueError(f"<UNK> embedding <UNK>{model_name}")
-
-    def get_dim(self):
-        return self.dim
-
-    def _get_embed_function(self, method: str):
-        """根据方法名返回对应的 embedding 函数"""
-        mapping = {
-            "openai": openai.openai_embed_sync,
-            "zhipu": zhipu.zhipu_embedding_sync,
-            "bedrock": bedrock.bedrock_embed_sync,
-            "hf": hf.hf_embed_sync,
-            "jina": jina.jina_embed_sync,
-            # "llama_index_impl": llama_index_impl.llama_index_embed,
-            "lollms": lollms.lollms_embed_sync,
-            "nvidia_openai": nvidia_openai.nvidia_openai_embed_sync,
-            "ollama": ollama.ollama_embed_sync,
-            "siliconcloud": siliconcloud.siliconcloud_embedding_sync,
-            "cohere": _cohere.cohere_embed_sync,
-            "mockembedder": lambda text, **kwargs: kwargs["embed_model"].encode(text).tolist(),
-            # "instructor": instructor.instructor_embed
-        }
-        if method not in mapping:
-            raise ValueError(f"不支持的 embedding 方法：{method}")
-
-        embed_fn = mapping[method]
-
-        return embed_fn
-
-    def _embed(self, text: str) -> list[float]:
-        """
-        异步执行 embedding 操作
-        :param text: 要 embedding 的文本
-        :param kwargs: embedding 方法可能需要的额外参数
-        :return: embedding 后的结果
-        """
-        return self.embed_fn(text, **self.kwargs)
-
-    def embed(self, text: str) -> list[float]:
-        return self._embed(text)
-    
-    def encode(self, text: str) -> list[float]:
-        return self._embed(text)
-    
-    @property
-    def method_name(self) -> str:
-        """当前embedding方法名"""
-        return self.init_method
-
-
-def apply_embedding_model(name: str = "default",**kwargs) -> EmbeddingModel:
-    """
-    usage  参见sage/api/model/operator_test.py
-    while name(method) = "hf", please set the param:model;
-    while name(method) = "openai",if you need call other APIs which are compatible with openai,set the params:base_url,api_key,model;
-    while name(method) = "jina/siliconcloud/cohere",please set the params:api_key,model;
-    Example:operator_test.py
-    """
-    return EmbeddingModel(method=name,**kwargs)
-
-def main():
-    embedding_model = EmbeddingModel(method="hf",model = "sentence-transformers/all-MiniLM-L6-v2")
-    for i in range(10):
-        start = time.time()
-        v = embedding_model.embed(f"{i} times ")
-        print(v)
-        end = time.time()
-        print(f"embedding time :{end-start}")
-        
-if __name__ =="__main__":
-    main()

sage/middleware/utils/embedding/hf.py

@@ -1,90 +0,0 @@
-import asyncio
-import copy
-import os
-from functools import lru_cache
-import asyncio
-import pipmaster as pm  # Pipmaster for dynamic library install
-
-# Dependencies should be installed via requirements.txt
-# transformers, torch, tenacity, and numpy are required for this module
-
-try:
-    from transformers import AutoTokenizer, AutoModelForCausalLM, AutoModel
-except ImportError:
-    raise ImportError(
-        "transformers package is required for HuggingFace embedding functionality. "
-        "Please install it via: pip install transformers"
-    )
-
-try:
-    import torch
-except ImportError:
-    raise ImportError(
-        "torch package is required for HuggingFace embedding functionality. "
-        "Please install it via: pip install torch"
-    )
-
-try:
-    import tenacity
-except ImportError:
-    raise ImportError(
-        "tenacity package is required for HuggingFace embedding functionality. "
-        "Please install it via: pip install tenacity"
-    )
-
-try:
-    import numpy
-except ImportError:
-    raise ImportError(
-        "numpy package is required for HuggingFace embedding functionality. "
-        "Please install it via: pip install numpy"
-    )
-
-os.environ["TOKENIZERS_PARALLELISM"] = "false"
-
-@lru_cache(maxsize=1)
-def initialize_hf_model(model_name):
-    hf_tokenizer = AutoTokenizer.from_pretrained(
-        model_name, device_map="auto", trust_remote_code=True
-    )
-    hf_model = AutoModelForCausalLM.from_pretrained(
-        model_name, device_map="auto", trust_remote_code=True
-    )
-    if hf_tokenizer.pad_token is None:
-        hf_tokenizer.pad_token = hf_tokenizer.eos_token
-
-    return hf_model, hf_tokenizer
-
-
-import torch
-
-def hf_embed_sync(text: str, tokenizer, embed_model) -> list[float]:
-    """
-    使用 HuggingFace 模型同步生成文本 embedding。
-
-    Args:
-        text (str): 输入文本
-        tokenizer: 已加载的 tokenizer
-        embed_model: 已加载的 PyTorch embedding 模型
-
-    Returns:
-        list[float]: embedding 向量
-    """
-    device = next(embed_model.parameters()).device
-    encoded_texts = tokenizer(
-        text, return_tensors="pt", padding=True, truncation=True
-    ).to(device)
-
-    with torch.no_grad():
-        outputs = embed_model(
-            input_ids=encoded_texts["input_ids"],
-            attention_mask=encoded_texts["attention_mask"],
-        )
-        embeddings = outputs.last_hidden_state.mean(dim=1)
-
-    if embeddings.dtype == torch.bfloat16:
-        return embeddings.detach().to(torch.float32).cpu()[0].tolist()
-    else:
-        return embeddings.detach().cpu()[0].tolist()
-
-

sage/middleware/utils/embedding/instructor.py

@@ -1,10 +0,0 @@
-import numpy as np
-from InstructorEmbedding import INSTRUCTOR
-from sentence_transformers import SentenceTransformer
-
-
-async def instructor_embed(
-        texts: [str], model: str = "hkunlp/instructor-large"
-) -> [np.array]:
-    _model = SentenceTransformer(model)
-    return _model.encode(texts)

sage/middleware/utils/embedding/jina.py

@@ -1,115 +0,0 @@
-import asyncio
-import os
-import pipmaster as pm  # Pipmaster for dynamic library install
-
-# Dependencies should be installed via requirements.txt
-# tenacity is required for this module
-
-try:
-    import tenacity
-except ImportError:
-    raise ImportError(
-        "tenacity package is required for Jina embedding functionality. "
-        "Please install it via: pip install tenacity"
-    )
-
-try:
-    import aiohttp
-except ImportError:
-    raise ImportError(
-        "aiohttp package is required for Jina embedding functionality. "
-        "Please install it via: pip install aiohttp"
-    )
-
-import numpy as np
-
-
-async def fetch_data(url, headers, data):
-    async with aiohttp.ClientSession() as session:
-        async with session.post(url, headers=headers, json=data) as response:
-            response_json = await response.json()
-            data_list = response_json.get("data", [])
-            return data_list
-
-
-async def jina_embed(
-    text: str,
-    dimensions: int = 1024,
-    late_chunking: bool = False,
-    base_url: str = None,
-    api_key: str = None,
-    model:str="jina-embeddings-v3"
-) -> list[float]:
-    if api_key:
-        os.environ["JINA_API_KEY"] = api_key
-    url = "https://api.jina.ai/v1/embeddings" if not base_url else base_url
-    headers = {
-        "Content-Type": "application/json",
-        "Authorization": f"Bearer {os.environ['JINA_API_KEY']}",
-    }
-    data = {
-        "model": f"{model}",
-        "normalized": True,
-        "embedding_type": "float",
-        "dimensions": f"{dimensions}",
-        "late_chunking": late_chunking,
-        "input": text,
-    }
-    data_list = await fetch_data(url, headers, data)
-    print(data_list)
-    return data_list[0]['embedding']
-
-import os
-import requests
-
-def jina_embed_sync(
-    text: str,
-    dimensions: int = 1024,
-    late_chunking: bool = False,
-    base_url: str = None,
-    api_key: str = None,
-    model: str = "jina-embeddings-v3"
-) -> list[float]:
-    """
-    同步版本：调用 Jina AI embedding API 获取嵌入向量
-
-    Args:
-        text: 待嵌入的文本
-        dimensions: 嵌入维度
-        late_chunking: 是否开启 late chunking
-        base_url: 自定义 API 地址（可选）
-        api_key: Jina API 密钥
-        model: 使用的模型名
-
-    Returns:
-        list[float]: 嵌入向量
-    """
-    if api_key:
-        os.environ["JINA_API_KEY"] = api_key
-
-    url = base_url or "https://api.jina.ai/v1/embeddings"
-    headers = {
-        "Content-Type": "application/json",
-        "Authorization": f"Bearer {os.environ['JINA_API_KEY']}",
-    }
-    payload = {
-        "model": model,
-        "normalized": True,
-        "embedding_type": "float",
-        "dimensions": dimensions,
-        "late_chunking": late_chunking,
-        "input": text,
-    }
-
-    try:
-        response = requests.post(url, headers=headers, json=payload)
-        response.raise_for_status()
-        data = response.json()
-        return data["data"][0]["embedding"]
-    except Exception as e:
-        raise RuntimeError(f"Jina API call failed: {str(e)}")
-
-
-
-
-