33 lines
991 B
Python
33 lines
991 B
Python
import dashscope
|
|
from dashscope import TextEmbedding
|
|
from typing import List
|
|
|
|
|
|
class EmbeddingService:
|
|
def __init__(self):
|
|
from app.core.config import settings
|
|
self.model = settings.embedding_model
|
|
self.dimension = settings.embedding_dim
|
|
dashscope.api_key = settings.dashscope_api_key
|
|
|
|
def embed_texts(self, texts: List[str]) -> List[List[float]]:
|
|
"""批量文本嵌入"""
|
|
response = TextEmbedding.call(
|
|
model=self.model,
|
|
input=texts,
|
|
)
|
|
|
|
if response.status_code == 200:
|
|
embeddings = []
|
|
for item in response.output.embeddings:
|
|
embeddings.append(item.embedding)
|
|
return embeddings
|
|
raise Exception(f"Embedding failed: {response.code}")
|
|
|
|
def embed_single(self, text: str) -> List[float]:
|
|
"""单个文本嵌入"""
|
|
embeddings = self.embed_texts([text])
|
|
return embeddings[0]
|
|
|
|
|
|
embedding_service = EmbeddingService() |