update sth. at 2025-01-16 21:48:33

app.py

@@ -1,32 +1,54 @@
-import gradio as gr
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel
 from transformers import AutoTokenizer, AutoModel
 import torch
+from typing import List, Dict
+import uvicorn
 
-# 1. 加载模型和分词器
-model_name = "jinaai/jina-embeddings-v3"  # 替换为您实际使用的模型名
+# 定义请求和响应模型
+class TextRequest(BaseModel):
+    text: str
+
+class EmbeddingResponse(BaseModel):
+    status: str
+    embeddings: List[List[float]]
+
+# 创建FastAPI应用
+app = FastAPI(
+    title="Jina Embeddings API",
+    description="Text embedding generation service using jina-embeddings-v3",
+    version="1.0.0"
+)
+
+# 加载模型和分词器
+model_name = "jinaai/jina-embeddings-v3"
 tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
 model = AutoModel.from_pretrained(model_name, trust_remote_code=True)
 
-# 2. 定义生成嵌入的函数
-def generate_embeddings(text):
-    # 使用分词器处理输入文本
-    inputs = tokenizer(text, return_tensors="pt")
-    # 禁用梯度计算，以减少资源消耗
-    with torch.no_grad():
-        # 获取最后一层隐藏状态并计算平均值作为嵌入
-        embeddings = model(**inputs).last_hidden_state.mean(dim=1)
-    # 将嵌入转换为Python列表，方便Gradio输出
-    return embeddings.numpy().tolist()
+@app.post("/generate_embeddings", response_model=EmbeddingResponse)
+async def generate_embeddings(request: TextRequest):
+    try:
+        # 使用分词器处理输入文本
+        inputs = tokenizer(request.text, return_tensors="pt", truncation=True, max_length=512)
+        
+        # 生成嵌入
+        with torch.no_grad():
+            embeddings = model(**inputs).last_hidden_state.mean(dim=1)
+        
+        return EmbeddingResponse(
+            status="success",
+            embeddings=embeddings.numpy().tolist()
+        )
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
 
-# 3. 使用Gradio定义接口
-iface = gr.Interface(
-    fn=generate_embeddings,  # 调用嵌入生成函数
-    inputs="text",           # 输入类型为文本
-    outputs="json",          # 输出为JSON格式，方便API调用
-    title="Text Embedding Generator",
-    description="Enter text to generate embeddings using the Jina model."
-)
+@app.get("/")
+async def root():
+    return {
+        "status": "active",
+        "model": model_name,
+        "usage": "Send POST request to /generate_embeddings"
+    }
 
-# 4. 启动Gradio应用
 if __name__ == "__main__":
-    iface.launch()
+    uvicorn.run(app, host="0.0.0.0", port=7860)

requirements.txt

@@ -1,3 +1,6 @@
 transformers
 torch
-einops
+einops
+fastapi
+uvicorn
+pydantic