ARQ-RAG / shared /embed.py
neshaki091
Feature: Switch to local sentence-transformers embedding for full independence
743476a
import logging
from sentence_transformers import SentenceTransformer
logger = logging.getLogger("Embedding-HuggingFace")
class EmbeddingManager:
"""
EmbeddingManager phiên bản Hugging Face TỰ VẬN HÀNH.
Sử dụng sentence-transformers để chạy trực tiếp trên CPU của Space.
"""
def __init__(self, model_name="nomic-ai/nomic-embed-text-v1.5"):
logger.info(f"🚀 Đang khởi tạo mô hình nhúng local: {model_name}")
# Chạy trên CPU để tiết kiệm tài nguyên và tương thích tốt trên HF Free
self.model = SentenceTransformer(model_name, trust_remote_code=True)
logger.info("✅ Khởi tạo mô hình nhúng thành công.")
def get_embedding(self, text):
try:
# Prefix 'search_query: ' là bắt buộc đối với mô hình Nomic khi dùng cho RAG retrieval
query_text = f"search_query: {text}"
embedding = self.model.encode(query_text)
return embedding.tolist()
except Exception as e:
logger.error(f"Lỗi khi nhúng câu hỏi tại local: {e}")
raise e