scripts/embed_data.py

import os
import sys
import time
import argparse

# Set UTF-8 encoding cho console
os.environ['PYTHONIOENCODING'] = 'utf-8'

# Thêm thư mục cha vào sys.path để import các module từ thư mục backend
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from core.data_processor import DataProcessor
from core.embedding_model import get_embedding_model

# Cấu hình logging phù hợp với environment
def setup_logging():
    import logging
    
    if os.getenv("SPACE_ID"):
        # Trên HuggingFace, chỉ log ra console
        logging.basicConfig(
            level=logging.INFO,
            format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
            handlers=[logging.StreamHandler()]
        )
    else:
        # Local environment
        logging.basicConfig(
            level=logging.INFO,
            format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
            handlers=[
                logging.StreamHandler(),
                logging.FileHandler("embed_data.log", encoding='utf-8')
            ]
        )
    
    return logging.getLogger("embed_data")

logger = setup_logging()

def embed_all_data(data_dir, force=False):
    """
    Embedding tất cả dữ liệu từ thư mục data
    """
    logger.info(f"Bat dau qua trinh embedding du lieu tu {data_dir}")
    start_time = time.time()
    
    try:
        # Khởi tạo các components
        logger.info("Khoi tao data processor...")
        data_processor = DataProcessor(data_dir=data_dir)
        
        logger.info("Khoi tao embedding model...")
        embedding_model = get_embedding_model()
        
        # Kiểm tra xem có chỉ mục hiện có không
        collection_size = embedding_model.count()
        logger.info(f"Kich thuoc collection hien tai: {collection_size}")
        
        if collection_size > 0 and not force:
            logger.info(f"Da ton tai chi muc voi {collection_size} items")
            end_time = time.time()
            logger.info(f"Hoan thanh kiem tra chi muc trong {end_time - start_time:.2f} giay")
            return True
        
        # Nếu buộc tạo lại hoặc chưa có chỉ mục, tạo mới
        if force and collection_size > 0:
            logger.info("Xoa chi muc cu va tao lai...")
            try:
                embedding_model.delete_collection()
                logger.info("Da xoa va tao lai collection")
            except Exception as e:
                logger.error(f"Loi khi xoa collection: {e}")
        
        # Chuẩn bị dữ liệu cho embedding
        logger.info("Dang chuan bi du lieu cho qua trinh embedding...")
        all_items = data_processor.prepare_for_embedding()
        logger.info(f"Da chuan bi {len(all_items)} items de embedding")
        
        if not all_items:
            logger.warning("Khong co items nao de embedding")
            return False
        
        # Thống kê các loại dữ liệu
        text_chunks = len([item for item in all_items if item.get("metadata", {}).get("content_type") == "text"])
        tables = len([item for item in all_items if item.get("metadata", {}).get("content_type") == "table"])
        figures = len([item for item in all_items if item.get("metadata", {}).get("content_type") == "figure"])
        
        logger.info(f"Bao gom: {text_chunks} van ban, {tables} bang bieu, {figures} hinh anh")
        
        # Thực hiện embedding
        logger.info("Bat dau qua trinh embedding...")
        success = embedding_model.index_chunks(all_items)
        
        if not success:
            logger.error("Loi xu ly embedding")
            return False
        
        end_time = time.time()
        elapsed_time = end_time - start_time
        
        # Kiểm tra kết quả cuối cùng
        final_count = embedding_model.count()
        logger.info(f"Hoan thanh qua trinh embedding {final_count} items trong {elapsed_time:.2f} giay")
        
        return True
        
    except Exception as e:
        logger.error(f"Loi embedding: {str(e)}")
        import traceback
        logger.error(traceback.format_exc())
        return False

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Embedding du lieu cho he thong Nutribot")
    parser.add_argument("--data-dir", type=str, default="data", 
                        help="Duong dan den thu muc chua du lieu (mac dinh: data)")
    parser.add_argument("--force", action="store_true", 
                        help="Xoa va tao lai chi muc neu da ton tai")
    
    args = parser.parse_args()
    
    # Chuẩn hóa đường dẫn
    data_dir = os.path.abspath(args.data_dir)
    
    # Kiểm tra thư mục data có tồn tại không
    if not os.path.exists(data_dir):
        logger.error(f"Thu muc {data_dir} khong ton tai!")
        sys.exit(1)
    
    # Thực hiện embedding
    success = embed_all_data(data_dir, args.force)
    if success:
        logger.info("Embedding hoan thanh thanh cong!")
        sys.exit(0)
    else:
        logger.error("Embedding that bai!")
        sys.exit(1)