update sth. at 2025-02-03 19:34:12

app.py

@@ -86,6 +86,7 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
+@app.post("/v1/embeddings", response_model=EmbeddingResponse)
 @app.post("/generate_embeddings", response_model=EmbeddingResponse)
 @app.post("/api/v1/embeddings", response_model=EmbeddingResponse)
 @app.post("/hf/v1/embeddings", response_model=EmbeddingResponse)

app.py——ok_baks

@@ -0,0 +1,146 @@
+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)

requirements.txt

@@ -9,4 +9,5 @@ numpy
 python-multipart
 sentencepiece
 safetensors
-
+pydantic
+click