app.py

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 uvicorn
import numpy as np

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.setup_logging()
        # CPU优化
        torch.set_num_threads(4)

    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

    async def _generate_embedding_internal(self, text: str) -> List[float]:
        """内部嵌入生成函数"""
        if not text.strip():
            raise ValueError("输入文本不能为空")

        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]

    @lru_cache(maxsize=1000)
    def get_cached_embedding(self, text: str) -> List[float]:
        """缓存包装函数"""
        loop = asyncio.new_event_loop()
        try:
            return loop.run_until_complete(self._generate_embedding_internal(text))
        finally:
            loop.close()

# 初始化服务
embedding_service = EmbeddingService()
app = FastAPI(
    title="Jina Embeddings API",
    description="Text embedding generation service using jina-embeddings-v3",
    version="1.0.0"
)

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:
        embedding = embedding_service.get_cached_embedding(request.input)
        return EmbeddingResponse(
            status="success",
            embeddings=[embedding]
        )
    except ValueError as e:
        raise HTTPException(status_code=400, detail=str(e))
    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)
    }

# Gradio界面
def gradio_interface(text: str) -> Dict:
    try:
        embedding = embedding_service.get_cached_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
    )