임베딩에서 통찰력까지: TwelveLabs Marengo 및 S3 벡터를 활용한 실전 크로스 모달 검색

import boto3
import json
import numpy as np
import uuid
import time
import os
from typing import List, Dict
from botocore.exceptions import ClientError

# AWS Configuration
AWS_REGION = "us-east-1"  # Update to your region
AWS_PROFILE = "default"   # Update to your profile

# S3 Vectors Configuration
VECTOR_BUCKET_NAME = "marengo-vectors-" + str(uuid.uuid4())[:8]
VECTOR_INDEX_NAME = "embeddings-index"
VECTOR_DIMENSION = 1024  # Marengo embedding dimension

# Marengo Model
MODEL_ID = 'twelvelabs.marengo-embed-2-7-v1:0'

# Temporary S3 bucket for Bedrock output (required by Bedrock API)
TEMP_S3_BUCKET = "<YOUR_S3_BUCKET>"  # TODO: Replace with your S3 bucket name

print(f"Vector Bucket: {VECTOR_BUCKET_NAME}")
print(f"Vector Index: {VECTOR_INDEX_NAME}")
print(f"Model: {MODEL_ID}")

# Initialize AWS clients
session = boto3.Session(profile_name=AWS_PROFILE)
bedrock_client = session.client('bedrock-runtime', region_name=AWS_REGION)
s3_client = session.client('s3')
s3vectors_client = session.client('s3vectors', region_name=AWS_REGION)
print("AWS clients initialized")

# Create Vector Bucket
try:
    s3vectors_client.create_vector_bucket(
        vectorBucketName=VECTOR_BUCKET_NAME,
        encryptionConfiguration={'sseType': 'AES256'}
    )
    print(f"Vector bucket '{VECTOR_BUCKET_NAME}' created")
except ClientError as e:
    if e.response['Error']['Code'] == 'ConflictException':
        print(f"Vector bucket already exists")
    else:
        print(f"Error: {e}")

# Create Vector Index
try:
    s3vectors_client.create_index(
        vectorBucketName=VECTOR_BUCKET_NAME,
        indexName=VECTOR_INDEX_NAME,
        dataType='float32',
        dimension=VECTOR_DIMENSION,
        distanceMetric='cosine',
        metadataConfiguration={'nonFilterableMetadataKeys': ['source']}
    )
    print(f"Index '{VECTOR_INDEX_NAME}' created")
except ClientError as e:
    if e.response['Error']['Code'] == 'ConflictException':
        print(f"Index already exists")
    else:
        print(f"Error: {e}")

def upload_file_to_s3(local_path: str, bucket: str, key: str) -> str:
    """
    Upload a local file to S3
    
    Args:
        local_path: Path to local file
        bucket: S3 bucket name
        key: S3 object key (path in bucket)
    
    Returns:
        S3 URI of uploaded file
    """
    try:
        s3_client.upload_file(local_path, bucket, key)
        s3_uri = f"s3://{bucket}/{key}"
        print(f"✅ Uploaded {os.path.basename(local_path)} to {s3_uri}")
        return s3_uri
    except ClientError as e:
        print(f"❌ Error uploading file: {e}")
        raise

def generate_text_embedding(text: str) -> List[float]:
    """
    Generate embedding for text using Marengo on Bedrock
    """
    # Create unique output path
    output_prefix = f'embeddings/{uuid.uuid4()}'
    
    # Start async embedding generation
    response = bedrock_client.start_async_invoke(
        modelId=MODEL_ID,
        modelInput={
            "inputType": "text",
            "inputText": text
        },
        outputDataConfig={
            "s3OutputDataConfig": {
                "s3Uri": f's3://{TEMP_S3_BUCKET}/{output_prefix}'
            }
        }
    )
    
    invocation_arn = response["invocationArn"]
    print(f"Generating text embedding for: '{text[:50]}...'")
    
    # Wait for completion
    status = None
    while status not in ["Completed", "Failed"]:
        response = bedrock_client.get_async_invoke(invocationArn=invocation_arn)
        status = response['status']
        time.sleep(2)
    
    if status != "Completed":
        raise Exception(f"Embedding generation failed")
    
    # Retrieve embedding from S3
    response = s3_client.list_objects_v2(Bucket=TEMP_S3_BUCKET, Prefix=output_prefix)
    for obj in response.get('Contents', []):
        if obj['Key'].endswith('output.json'):
            result = s3_client.get_object(Bucket=TEMP_S3_BUCKET, Key=obj['Key'])
            data = json.loads(result['Body'].read())
            return data['data'][0]['embedding']
    
    raise Exception("No embedding output found")

def generate_video_embedding(video_path: str, start_sec: float = 0, length_sec: float = None) -> List[float]:
    """
    Generate embedding for video using Marengo on Bedrock
    """
    # Read video file and encode to base64
    with open(video_path, 'rb') as video_file:
        video_base64 = base64.b64encode(video_file.read()).decode('utf-8')
    
    # Build model input with optional time segments
    model_input = {
        "inputType": "video",
        "mediaSource": {"base64String": video_base64},
        "embeddingOption": ["visual-text", "audio"]  # Capture both visual and audio
    }
    
    if start_sec is not None:
        model_input["startSec"] = start_sec
    if length_sec is not None:
        model_input["lengthSec"] = length_sec
    
    # ... rest follows the same async pattern

def generate_video_embedding_from_s3(s3_uri: str, start_sec: float = 0, 
                                    length_sec: float = None) -> List[float]:
    """Generate embedding for video from S3 using Marengo on Bedrock"""    output_prefix = f'{OUTPUT_PREFIX}video/{uuid.uuid4()}'
    
    model_input = {
        "inputType": "video",
        "mediaSource": {
            "s3Location": {
                "uri": s3_uri,
                "bucketOwner": ACCOUNT_ID
            }
        },
        "embeddingOption": ["visual-text", "audio"]
    }
    
    response = bedrock_client.start_async_invoke(
        modelId=MODEL_ID,
        modelInput=model_input,
        outputDataConfig={
            "s3OutputDataConfig": {
                "s3Uri": f's3://{TEMP_S3_BUCKET}/{output_prefix}',
                "bucketOwner": ACCOUNT_ID
            }
        }
    )
    
    print(f"🎬 Generating video embedding from S3: {s3_uri}")
    # Enhanced status monitoring with retry logic...

# Upload video to S3
video_s3_key = f"{MEDIA_PREFIX}videos/sample_video.mp4"
video_s3_uri = upload_file_to_s3('video.mp4', TEMP_S3_BUCKET, video_s3_key)

# Generate embedding from S3
video_embedding_s3 = generate_video_embedding_from_s3(video_s3_uri, start_sec=0, length_sec=10)

# Compare VIDEO embeddings
print("🎬 VIDEO EMBEDDING VERIFICATION")
print("-" * 50)

# Generate using Base64
print("Generating via Base64...")
video_base64_emb = generate_video_embedding('video.mp4', start_sec=0, length_sec=10)

# Upload and generate using S3 URI
print("Generating via S3 URI...")
video_s3_key = f"{MEDIA_PREFIX}verification/video.mp4"
video_s3_uri = upload_file_to_s3('video.mp4', TEMP_S3_BUCKET, video_s3_key)
video_s3_emb = generate_video_embedding_from_s3(video_s3_uri, start_sec=0, length_sec=10)

# Calculate similarity
similarity = cosine_similarity(video_base64_emb, video_s3_emb)
print(f"📊 Results:")
print(f"  Cosine similarity: {similarity:.6f}")
print(f"  Are they identical? {'✅ YES' if similarity > 0.9999 else '❌ NO'}")

# Visualize text embeddings
if text_embeddings:
    plt.figure(figsize=(12, 3))
    for i, (text, emb, _) in enumerate(text_embeddings[:3]):
        plt.subplot(1, 3, i+1)
        plt.imshow([emb[:100]], aspect='auto', cmap='coolwarm')
        plt.title(f'Text {i+1}', fontsize=10)
        plt.xlabel('Dimensions (first 100)')
        plt.colorbar(orientation='horizontal', pad=0.1)
    plt.suptitle('Text Embeddings Visualization', fontsize=12)
    plt.tight_layout()
    plt.show()

def store_embeddings(embeddings_data: List[tuple]) -> bool:
    """
    Store embeddings in S3 Vectors with metadata
    """
    vectors_to_insert = []
    
    for i, item in enumerate(embeddings_data):
        # Support (text, embedding) and (text, embedding, media_type)
        if len(item) == 3:
            text, embedding, media_type = item
        else:
            text, embedding = item
            media_type = "text"
        
        vector_entry = {
            'key': f'vector_{i:04d}',           # deterministic, idempotent upserts
            'data': {
                'float32': [float(v) for v in embedding]  # convert to float32-compatible list
            },
            'metadata': {
                'text': text,                   # human-readable description
                'media_type': media_type,       # useful for UI badges and filtering
                'id': i                         # numeric id for quick reference
            }
        }
        vectors_to_insert.append(vector_entry)

    try:
        s3vectors_client.put_vectors(
            vectorBucketName=VECTOR_BUCKET_NAME,
            indexName=VECTOR_INDEX_NAME,
            vectors=vectors_to_insert
        )
        print(f"Stored {len(vectors_to_insert)} vectors in S3 Vectors")
        
        # Optional: quick summary
        media_counts = {}
        for item in embeddings_data:
            mtype = item[2] if len(item) == 3 else "text"
            media_counts[mtype] = media_counts.get(mtype, 0) + 1
        print("Media types stored:")
        for m, c in media_counts.items():
            print(f"  - {m}: {c}")
        return True
    except ClientError as e:
        print(f"Error storing vectors: {e}")
        return False

def search_similar(query_text: str = None, query_embedding: List[float] = None, top_k: int = 3) -> List[Dict]:
    """
    Search for similar vectors using either a text query or a pre-computed embedding
    """
    # Generate embedding for query if text is provided
    if query_text and not query_embedding:
        print(f"\nSearching for: '{query_text}'")
        query_embedding = generate_text_embedding(query_text)
    elif query_embedding:
        print(f"\nSearching with provided embedding...")
    else:
        raise ValueError("Either query_text or query_embedding must be provided")
    
    # Query S3 Vectors
    try:
        response = s3vectors_client.query_vectors(
            vectorBucketName=VECTOR_BUCKET_NAME,
            indexName=VECTOR_INDEX_NAME,
            topK=top_k,
            queryVector={'float32': [float(v) for v in query_embedding]},
            returnMetadata=True,
            returnDistance=True
        )
        
        results = []
        for vector in response.get('vectors', []):
            # Handle both old vectors (without media_type) and new vectors (with media_type)
            results.append({
                'text': vector['metadata'].get('text', 'No description'),
                'media_type': vector['metadata'].get('media_type', 'text'),  # Default to 'text' for old vectors
                'similarity': 1.0 - vector.get('distance', 0),  # Convert distance to similarity
                'key': vector.get('key')
            })
        
        return results
        
    except ClientError as e:
        print(f"Query failed: {e}")
        return []

# Example queries
queries = [
    "Someone enjoying the ocean view",
    "People talking",
    "Cooking activities"
]

for query in queries:
    results = search_similar(query, top_k=2)
    print(f"\nResults for: '{query}'")
    for i, result in enumerate(results, 1):
        print(f"  {i}. Similarity: {result['similarity']:.3f} | Text: {result['text']}")

if all_results:
    fig, axes = plt.subplots(len(all_results), 1, figsize=(10, len(all_results)*2))
    if len(all_results) == 1:
        axes = [axes]
    
    for idx, result_set in enumerate(all_results):
        query = result_set['query']
        results = result_set['results']
        
        if results:
            similarities = [r['similarity'] for r in results]
            labels = [f"{r['media_type'][:3]}" for r in results]
            colors = ['#4CAF50' if r['media_type'] == 'text' else 
                     '#2196F3' if r['media_type'] == 'video' else
                     '#FF9800' if r['media_type'] == 'audio' else
                     '#9C27B0' for r in results]
            
            axes[idx].barh(range(len(similarities)), similarities, color=colors)
            axes[idx].set_yticks(range(len(similarities)))
            axes[idx].set_yticklabels(labels)
            axes[idx].set_xlabel('Similarity Score')
            axes[idx].set_title(f'Query: "{query[:40]}..."')
            axes[idx].set_xlim(0, 1)
            
            # Add value labels
            for i, v in enumerate(similarities):
                axes[idx].text(v, i, f' {v:.3f}', va='center')
    
    plt.suptitle('Cross-Modal Search Results', fontsize=14, fontweight='bold')
    plt.tight_layout()
    plt.savefig('cross_modal_search_results.png', dpi=150, bbox_inches='tight')
    plt.show()

print("\n💾 Search results saved to JSON files")

import numpy as np

def cosine_similarity(vec1, vec2) -> float:
    """
    Calculate cosine similarity between two vectors.
    """
    v1 = np.array(vec1, dtype=np.float32)
    v2 = np.array(vec2, dtype=np.float32)
    dot = np.dot(v1, v2)
    n1 = np.linalg.norm(v1)
    n2 = np.linalg.norm(v2)
    return float(dot / (n1 * n2))

import matplotlib.pyplot as plt

# Example texts
text1 = "A beautiful sunset over the ocean"
text2 = "A sunrise at the beach"
text3 = "People working in an office"

print("Generating embeddings for comparison...")
emb1 = generate_text_embedding(text1)   # from previous section
emb2 = generate_text_embedding(text2)
emb3 = generate_text_embedding(text3)

sim_12 = cosine_similarity(emb1, emb2)
sim_13 = cosine_similarity(emb1, emb3)
sim_23 = cosine_similarity(emb2, emb3)

print("\nDirect Similarity Comparison:")
print(f"'{text1}' vs '{text2}': {sim_12:.3f}")
print(f"'{text1}' vs '{text3}': {sim_13:.3f}")
print(f"'{text2}' vs '{text3}': {sim_23:.3f}")

labels = ['Sunset/Ocean', 'Sunrise/Beach', 'Office Work']

# Assemble a 3x3 similarity matrix
similarity_matrix = np.array([
    [1.0,    sim_12, sim_13],
    [sim_12, 1.0,    sim_23],
    [sim_13, sim_23, 1.0]
], dtype=np.float32)

plt.figure(figsize=(8, 6))
im = plt.imshow(similarity_matrix, cmap='RdYlGn', vmin=0, vmax=1, aspect='auto')
plt.colorbar(im, label='Cosine Similarity')

# Annotations
for i in range(similarity_matrix.shape[0]):
    for j in range(similarity_matrix.shape[1]):
        plt.text(j, i, f'{similarity_matrix[i, j]:.3f}',
                 ha="center", va="center", color="black", fontweight='bold')

plt.xticks(range(3), labels, rotation=45, ha='right')
plt.yticks(range(3), labels)
plt.title('Embedding Similarity Matrix', fontsize=14, fontweight='bold')
plt.tight_layout()
plt.savefig('similarity_matrix.png', dpi=150, bbox_inches='tight')
plt.show()