Spaces:
Running
Running
File size: 4,557 Bytes
8e90b2d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 |
import asyncio
import logging
import torch
import gradio as gr
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from typing import List, Dict
from functools import lru_cache
import numpy as np
from threading import Lock
import uvicorn
class EmbeddingRequest(BaseModel):
input: str
model: str = "jinaai/jina-embeddings-v3"
class EmbeddingResponse(BaseModel):
status: str
embeddings: List[List[float]]
class EmbeddingService:
def __init__(self):
self.model_name = "jinaai/jina-embeddings-v3"
self.max_length = 512
self.device = torch.device("cpu")
self.model = None
self.tokenizer = None
self.lock = Lock()
self.setup_logging()
torch.set_num_threads(4) # CPU优化
def setup_logging(self):
logging.basicConfig(
level=logging.INFO,
format='%(asctime)s - %(levelname)s - %(message)s'
)
self.logger = logging.getLogger(__name__)
async def initialize(self):
try:
from transformers import AutoTokenizer, AutoModel
self.tokenizer = AutoTokenizer.from_pretrained(
self.model_name,
trust_remote_code=True
)
self.model = AutoModel.from_pretrained(
self.model_name,
trust_remote_code=True
).to(self.device)
self.model.eval()
torch.set_grad_enabled(False)
self.logger.info(f"模型加载成功，使用设备: {self.device}")
except Exception as e:
self.logger.error(f"模型初始化失败: {str(e)}")
raise
@lru_cache(maxsize=1000)
def get_embedding(self, text: str) -> List[float]:
"""同步生成嵌入向量，带缓存"""
with self.lock:
try:
inputs = self.tokenizer(
text,
return_tensors="pt",
truncation=True,
max_length=self.max_length,
padding=True
)
with torch.no_grad():
outputs = self.model(**inputs).last_hidden_state.mean(dim=1)
return outputs.numpy().tolist()[0]
except Exception as e:
self.logger.error(f"生成嵌入向量失败: {str(e)}")
raise
embedding_service = EmbeddingService()
app = FastAPI()
app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
@app.post("/generate_embeddings", response_model=EmbeddingResponse)
@app.post("/api/v1/embeddings", response_model=EmbeddingResponse)
@app.post("/hf/v1/embeddings", response_model=EmbeddingResponse)
@app.post("/api/v1/chat/completions", response_model=EmbeddingResponse)
@app.post("/hf/v1/chat/completions", response_model=EmbeddingResponse)
async def generate_embeddings(request: EmbeddingRequest):
try:
# 使用run_in_executor避免事件循环问题
embedding = await asyncio.get_running_loop().run_in_executor(
None,
embedding_service.get_embedding,
request.input
)
return EmbeddingResponse(
status="success",
embeddings=[embedding]
)
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
@app.get("/")
async def root():
return {
"status": "active",
"model": embedding_service.model_name,
"device": str(embedding_service.device)
}
def gradio_interface(text: str) -> Dict:
try:
embedding = embedding_service.get_embedding(text)
return {
"status": "success",
"embeddings": [embedding]
}
except Exception as e:
return {
"status": "error",
"message": str(e)
}
iface = gr.Interface(
fn=gradio_interface,
inputs=gr.Textbox(lines=3, label="输入文本"),
outputs=gr.JSON(label="嵌入向量结果"),
title="Jina Embeddings V3",
description="使用jina-embeddings-v3模型生成文本嵌入向量",
examples=[["这是一个测试句子。"]]
)
@app.on_event("startup")
async def startup_event():
await embedding_service.initialize()
if __name__ == "__main__":
asyncio.run(embedding_service.initialize())
gr.mount_gradio_app(app, iface, path="/ui")
uvicorn.run(app, host="0.0.0.0", port=7860, workers=1)
|