File size: 558 Bytes
c69a4d6 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 |
#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time : 2025/4/24 11:50
# @Author : hukangzhe
# @File : embedder.py
# @Description :
from sentence_transformers import SentenceTransformer
from typing import List
import numpy as np
class EmbeddingModel:
def __init__(self, model_name: str):
self.embedding_model = SentenceTransformer(model_name)
def embed(self, texts: List[str], batch_size: int = 32) -> np.ndarray:
return self.embedding_model.encode(texts, batch_size=batch_size, convert_to_numpy=True)
|