TwelveLabs Embed API Beta

# 필요한 라이브러리 및 의존성 설치
!pip install twelvelabs
# 필요한 모듈 임포트
import twelvelabs

from twelvelabs import TwelveLabs
from twelvelabs.models.embed import EmbeddingsTask

# Twelve Labs 클라이언트 초기화
twelvelabs_client = TwelveLabs(api_key=TL_API_KEY)

def generate_embedding(video_url):
    # 임베딩 작업 생성
    task = twelvelabs_client.embed.task.create(
        engine_name="Marengo-retrieval-2.6",
        video_url=video_url
    )
    print(f"태스크 생성 완료: id={task.id} engine_name={task.engine_name} status={task.status}")

    # 작업 진행 상황을 모니터링할 콜백 함수 정의
    def on_task_update(task: EmbeddingsTask):
        print(f"  상태={task.status}")

    # 작업이 완료될 때까지 대기
    status = task.wait_for_done(
        sleep_interval=2,
        callback=on_task_update
    )
    print(f"임베딩 진행 완료: {status}")

    # 작업 결과 가져오기
    task_result = twelvelabs_client.embed.task.retrieve(task.id)

    # 임베딩 추출 및 반환
    embeddings = []
    for v in task_result.video_embeddings:
        embeddings.append({
            'embedding': v.values,
            'start_offset_sec': v.start_offset_sec,
            'end_offset_sec': v.end_offset_sec,
            'embedding_scope': v.embedding_scope
        })

    return embeddings, task_result

# 사용 예시
video_url = "https://storage.googleapis.com/ad-demos-datasets/videos/Ecommerce%20v2.5.mp4"

# 비디오에 대한 임베딩 생성
embeddings, task_result = generate_embedding(video_url)

print(f"비디오에서 {len(embeddings)}개의 임베딩이 생성되었습니다")
for i, emb in enumerate(embeddings):
    print(f"임베딩 {i+1}:")
    print(f"  범위: {emb['embedding_scope']}")
    print(f"  시간 범위: {emb['start_offset_sec']} - {emb['end_offset_sec']} 초")
    print(f"  임베딩 벡터 (앞의 5개 값): {emb['embedding'][:5]}")
    print()

def create_text_embedding(
    twelvelabs_client: TwelveLabs,
    text: str,
    engine_name: str = "Marengo-retrieval-2.6",
    verbose: bool = True
) -> dict:
    """
    Twelve Labs Embed API를 사용해 텍스트 임베딩을 생성합니다.
    
    예시:
        twelvelabs_client = TwelveLabs(api_key=TL_API_KEY)
        text_embedding = create_text_embedding(twelvelabs_client, "여기에 텍스트 입력")
    """
    # 임베딩 생성
    text_embedding = twelvelabs_client.embed.create(
        engine_name=engine_name,
        text=text
    )
    
    # verbose 속성이 True인 경우 정보 출력
    if verbose:
        print("텍스트 임베딩이 생성되었습니다")
        print(f" 엔진: {text_embedding.engine_name}")
        print(f" 임베딩: {text_embedding.text_embedding}")
    
    return text_embedding

def create_audio_embedding(
    twelvelabs_client: TwelveLabs,
    audio_file: Union[str, Path],
    engine_name: str = "Marengo-retrieval-2.6",
    verbose: bool = True
) -> dict:
    """
    Twelve Labs Embed API를 사용해 오디오 임베딩을 생성합니다.
    
    예시:
        twelvelabs_client = TwelveLabs(api_key=TL_API_KEY)
        audio_embedding = create_audio_embedding(twelvelabs_client, "path/to/audio.mp3")
    """
    # 임베딩 생성
    audio_embedding = twelvelabs_client.embed.create(
        engine_name=engine_name,
        audio_file=audio_file,
    )
    
    # verbose 속성이 True인 경우 정보 출력
    if verbose:
        print("오디오 임베딩이 생성되었습니다")
        print(f" 엔진: {audio_embedding.engine_name}")
        
        if audio_embedding.audio_embedding.segments:
            print("세그먼트 목록:")
            for i, segment in enumerate(audio_embedding.audio_embedding.segments, 1):
                print(f" 세그먼트 {i}:")
                print(f" 시작 오프셋 (초): {segment.start_offset_sec}")
                if segment.values:
                    print(f" 값: {segment.values[:5]}... (일부 생략)")
    
    return audio_embedding

def create_image_embedding(
    twelvelabs_client: TwelveLabs,
    image_file: Union[str, Path],
    engine_name: str = "Marengo-retrieval-2.6",
    verbose: bool = True
) -> dict:
    """
    Twelve Labs Embed API를 사용해 이미지 임베딩을 생성합니다.
    
    예시:
        twelvelabs_client = TwelveLabs(api_key=TL_API_KEY)
        image_embedding = create_image_embedding(twelvelabs_client, "path/to/image.jpg")
    """
    # 임베딩 생성
    image_embedding = twelvelabs_client.embed.create(
        engine_name=engine_name,
        image_file=image_file
    )
    
    # verbose 속성이 True인 경우 정보 출력
    if verbose:
        print("이미지 임베딩이 생성되었습니다")
        print(f" 엔진: {image_embedding.engine_name}")
        print(f" 임베딩: {image_embedding.image_embedding.values[:5]}... (일부 생략)")
    
    return image_embedding