172 lines
6.2 KiB
Python
172 lines
6.2 KiB
Python
from fastapi import FastAPI, Depends, HTTPException, status
|
||
from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
|
||
from sentence_transformers import SentenceTransformer
|
||
from pydantic import BaseModel
|
||
from fastapi.middleware.cors import CORSMiddleware
|
||
import uvicorn
|
||
import tiktoken
|
||
import numpy as np
|
||
from scipy.interpolate import interp1d
|
||
from typing import List, Optional
|
||
from sklearn.preprocessing import PolynomialFeatures
|
||
from sklearn.decomposition import PCA
|
||
import torch
|
||
import os
|
||
|
||
# 接口秘钥环境变量传入
|
||
sk_key = os.environ.get('sk-key', 'sk-aaabbbcccdddeeefffggghhhiiijjjkkk')
|
||
# 是否自动进行维度操作的环境变量，默认为false
|
||
auto_dim = os.environ.get('auto-dim', 'false').lower() == 'true'
|
||
# 模型名称, 必须在models文件夹下有对应的模型文件夹
|
||
model_name = os.environ.get('model-name', 'bge-m3')
|
||
|
||
# 创建一个FastAPI实例
|
||
app = FastAPI()
|
||
|
||
app.add_middleware(
|
||
CORSMiddleware,
|
||
allow_origins=["*"],
|
||
allow_credentials=True,
|
||
allow_methods=["*"],
|
||
allow_headers=["*"],
|
||
)
|
||
|
||
# 创建一个HTTPBearer实例
|
||
security = HTTPBearer()
|
||
|
||
# 预加载模型
|
||
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') # 检测是否有GPU可用，如果有则使用cuda设备，否则使用cpu设备
|
||
if torch.cuda.is_available():
|
||
print('本次加载模型的设备为GPU: ', torch.cuda.get_device_name(0))
|
||
else:
|
||
print('本次加载模型的设备为CPU.')
|
||
|
||
print(f'加载模型: {model_name}')
|
||
model = SentenceTransformer(f'./models/{model_name}',device=device)
|
||
|
||
# 创建PCA降维模型
|
||
pca = None
|
||
|
||
class EmbeddingRequest(BaseModel):
|
||
input: List[str]
|
||
model: str
|
||
dimensions: Optional[int] = 512
|
||
|
||
class EmbeddingResponse(BaseModel):
|
||
data: list
|
||
model: str
|
||
object: str
|
||
usage: dict
|
||
|
||
def num_tokens_from_string(string: str) -> int:
|
||
"""Returns the number of tokens in a text string."""
|
||
encoding = tiktoken.get_encoding('cl100k_base')
|
||
num_tokens = len(encoding.encode(string))
|
||
return num_tokens
|
||
|
||
# 插值法
|
||
def interpolate_vector(vector, target_length):
|
||
original_indices = np.arange(len(vector))
|
||
target_indices = np.linspace(0, len(vector)-1, target_length)
|
||
f = interp1d(original_indices, vector, kind='linear')
|
||
return f(target_indices)
|
||
|
||
def expand_features(embedding, target_length):
|
||
poly = PolynomialFeatures(degree=2)
|
||
expanded_embedding = poly.fit_transform(embedding.reshape(1, -1))
|
||
expanded_embedding = expanded_embedding.flatten()
|
||
if len(expanded_embedding) > target_length:
|
||
# 如果扩展后的特征超过目标长度，可以通过截断或其他方法来减少维度
|
||
expanded_embedding = expanded_embedding[:target_length]
|
||
elif len(expanded_embedding) < target_length:
|
||
# 如果扩展后的特征少于目标长度，可以通过填充或其他方法来增加维度
|
||
expanded_embedding = np.pad(expanded_embedding, (0, target_length - len(expanded_embedding)))
|
||
return expanded_embedding
|
||
|
||
# 降维方法：使用PCA将向量从1024维降到512维
|
||
def reduce_dimensions(embeddings, target_dim=512):
|
||
global pca
|
||
|
||
# 将列表转换为numpy数组
|
||
embeddings_array = np.array(embeddings)
|
||
|
||
# 检查样本数量
|
||
n_samples = embeddings_array.shape[0]
|
||
n_features = embeddings_array.shape[1]
|
||
|
||
# 如果只有一个样本，无法使用PCA，改用插值法
|
||
if n_samples == 1:
|
||
return [interpolate_vector(embeddings_array[0], target_dim)]
|
||
|
||
# 确保目标维度不超过可能的最大值
|
||
actual_target_dim = min(target_dim, n_samples, n_features)
|
||
if actual_target_dim < target_dim:
|
||
print(f"警告：目标维度{target_dim}超过了可能的最大值，已调整为{actual_target_dim}")
|
||
|
||
# 如果是第一次运行或者输入维度变化，重新初始化PCA
|
||
if pca is None or pca.n_components != actual_target_dim:
|
||
pca = PCA(n_components=actual_target_dim)
|
||
# 先拟合再转换
|
||
reduced_embeddings = pca.fit_transform(embeddings_array)
|
||
else:
|
||
# 直接使用已训练的PCA模型转换
|
||
reduced_embeddings = pca.transform(embeddings_array)
|
||
|
||
# 如果实际降维后的维度小于目标维度，使用插值法扩展
|
||
if actual_target_dim < target_dim:
|
||
reduced_embeddings = [interpolate_vector(embedding, target_dim) for embedding in reduced_embeddings]
|
||
|
||
return list(reduced_embeddings)
|
||
|
||
@app.post("/v1/embeddings", response_model=EmbeddingResponse)
|
||
async def get_embeddings(request: EmbeddingRequest, credentials: HTTPAuthorizationCredentials = Depends(security)):
|
||
|
||
if credentials.credentials != sk_key:
|
||
raise HTTPException(
|
||
status_code=status.HTTP_401_UNAUTHORIZED,
|
||
detail="Invalid authorization code",
|
||
)
|
||
|
||
# 计算嵌入向量和tokens数量
|
||
embeddings = [model.encode(text) for text in request.input]
|
||
|
||
# 检查是否需要进行维度操作
|
||
if auto_dim:
|
||
# 检查向量维度
|
||
embedding_dim = len(embeddings[0])
|
||
|
||
# 如果维度大于512，则降维到512
|
||
if embedding_dim > request.dimensions:
|
||
embeddings = reduce_dimensions(embeddings, target_dim=request.dimensions)
|
||
# 如果维度小于512，则使用插值法扩展到512
|
||
elif embedding_dim < request.dimensions:
|
||
embeddings = [interpolate_vector(embedding, request.dimensions) for embedding in embeddings]
|
||
|
||
# 归一化处理
|
||
embeddings = [embedding / np.linalg.norm(embedding) for embedding in embeddings]
|
||
# 将numpy数组转换为列表
|
||
embeddings = [embedding.tolist() for embedding in embeddings]
|
||
prompt_tokens = sum(len(text.split()) for text in request.input)
|
||
total_tokens = sum(num_tokens_from_string(text) for text in request.input)
|
||
|
||
response = {
|
||
"data": [
|
||
{
|
||
"embedding": embedding,
|
||
"index": index,
|
||
"object": "embedding"
|
||
} for index, embedding in enumerate(embeddings)
|
||
],
|
||
"model": model_name,
|
||
"object": "list",
|
||
"usage": {
|
||
"prompt_tokens": prompt_tokens,
|
||
"total_tokens": total_tokens,
|
||
}
|
||
}
|
||
|
||
return response
|
||
|
||
if __name__ == "__main__":
|
||
uvicorn.run("main:app", host='0.0.0.0', port=6008, workers=1)
|