TwelveLabsとVespaを活用したマルチベクトル動画検索

!pip3 install pyvespa vespacli twelvelabs pandas

import os
import hashlib
import json

from vespa.package import (
    ApplicationPackage,
    Field,
    Schema,
    Document,
    HNSW,
    RankProfile,
    FieldSet,
    SecondPhaseRanking,
    Function,
)

from vespa.deployment import VespaCloud
from vespa.io import VespaResponse, VespaQueryResponse

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

import pandas as pd

from datetime import datetime

TL_API_KEY = os.getenv("TL_API_KEY") or input("Enter your TL_API key: ")

# Replace with your tenant name from the Vespa Cloud Console
tenant_name = "vespa-team"
# Replace with your application name (does not need to exist yet)
application = "videosearch"

# Initialize Twelve Labs client
client = TwelveLabs(api_key=TL_API_KEY)

# Test the connection
try:
    client.tasks.list()
    print("Successfully connected to Twelve Labs API")
except Exception as e:
    print(f"Error connecting to Twelve Labs API: {e}")

VIDEO_URLs = [
    "https://ia801503.us.archive.org/27/items/hide-and-seek-with-giant-jenny/HnVideoEditor_2022_10_29_205557707.ia.mp4",
    "https://ia601401.us.archive.org/1/items/twas-the-night-before-christmas-1974-full-movie-freedownloadvideo.net/twas-the-night-before-christmas-1974-full-movie-freedownloadvideo.net.mp4",
    "https://dn720401.ca.archive.org/0/items/mr-bean-the-animated-series-holiday-for-teddy/S2E12.ia.mp4",
]

# Initialize client and create index
client = TwelveLabs(api_key=TL_API_KEY)
timestamp = int(datetime.now().timestamp())
index_name = "Vespa_" + str(timestamp)

# Create Index with Pegasus 1.2 model
index = client.index.create(
    name=index_name,
    models=[{"name": "pegasus1.2", "options": ["visual", "audio"]}],
    addons=["thumbnail"]
)

def on_task_update(task: EmbeddingsTask):
    print(f"  Status={task.status}")

for video_url in VIDEO_URLs:
    task = client.task.create(index_id=index.id, url=video_url)
    status = task.wait_for_done(sleep_interval=10, callback=on_task_update)
    if task.status != "ready":
        raise RuntimeError(f"Indexing failed with status {task.status}")

summaries = []
keywords_array = []
titles = [
    "Mr. Bean the Animated Series Holiday for Teddy",
    "Twas the night before Christmas",
    "Hide and Seek with Giant Jenny",
]

videos = client.index.video.list(index_id)
for video in videos:
    # Generate summary
    res = client.generate.summarize(
        video_id=video.id,
        type="summary",
        prompt="Generate an abstract of the video serving as metadata on the video, up to five sentences."
    )
    summaries.append(res.summary)
    
    # Generate keywords
    keywords = client.generate.text(
        video_id=video.id,
        prompt="Based on this video, I want to generate five keywords for SEO. Provide just the keywords as a comma delimited list."
    )
    keywords_array.append(keywords.data)

task_ids = []

for url in VIDEO_URLs:
    task = client.embed.task.create(model_name="Marengo-retrieval-2.7", video_url=url)
    task_ids.append(str(task.id))
    status = task.wait_for_done(sleep_interval=10, callback=on_task_update)
    if task.status != "ready":
        raise RuntimeError(f"Embedding failed with status {task.status}")

tasks = []
for task_id in task_ids:
    task = client.embed.task.retrieve(task_id)
    tasks.append(task)

# Define the schema for "videos"
videos_schema = Schema(
    name="videos",
    document=Document(
        fields=[
            Field(name="video_url", type="string", indexing=["summary"]),
            Field(name="title", type="string", indexing=["index", "summary"], match=["text"], index="enable-bm25"),
            Field(name="keywords", type="string", indexing=["index", "summary"], match=["text"], index="enable-bm25"),
            Field(name="video_summary", type="string", indexing=["index", "summary"], match=["text"], index="enable-bm25"),
            Field(name="embedding_scope", type="string", indexing=["attribute", "summary"]),
            Field(name="start_offset_sec", type="array<float>", indexing=["attribute", "summary"]),
            Field(name="end_offset_sec", type="array<float>", indexing=["attribute", "summary"]),
            Field(
                name="embeddings",
                type="tensor<float>(p{},x[1024])",
                indexing=["index", "attribute"],
                ann=HNSW(distance_metric="angular"),
            ),
        ]
    ),
)

# Add fieldsets for search
fieldsets = [
    FieldSet(name="default", fields=["title", "keywords", "video_summary"]),
]

# Define ranking functions
mapfunctions = [
    Function(
        name="similarities",
        expression="""sum(query(q) * attribute(embeddings), x)""",
    ),
    Function(
        name="bm25_score",
        expression="bm25(title) + bm25(keywords) + bm25(video_summary)",
    ),
]

# Define a hybrid rank profile
semantic_rankprofile = RankProfile(
    name="hybrid",
    inputs=[("query(q)", "tensor<float>(x[1024])")],
    first_phase="bm25_score",
    second_phase=SecondPhaseRanking(expression="closeness(field, embeddings)", rerank_count=10),
    match_features=["closest(embeddings)"],
    summary_features=["similarities"],
    functions=mapfunctions,
)

# Add rank profile to schema
videos_schema.add_rank_profile(semantic_rankprofile)

# Create the application package
package = ApplicationPackage(name=application, schema=[videos_schema])

from vespa.deployment import VespaCloud

# Deploy the application package to Vespa Cloud
vespa_cloud = VespaCloud(
    tenant=tenant_name,
    application=application,
    application_package=package,
    key_content=os.getenv("VESPA_TEAM_API_KEY"),  # Replace with your API key content
)

app = vespa_cloud.deploy()
print("Deployment complete!")

import hashlib

# Initialize a list to store Vespa feed documents
vespa_feed = []

# Reverse VIDEO_URLs since attributes were generated in reverse order
VIDEO_URLs.reverse()

for i, task in enumerate(tasks):
    video_url = VIDEO_URLs[i]
    title = titles[i]
    keywords = keywords_array[i]
    summary = summaries[i]

    start_offsets = []
    end_offsets = []
    embeddings = {}

    for index, segment in enumerate(task.video_embedding.segments):
        start_offsets.append(float(segment.start_offset_sec))
        end_offsets.append(float(segment.end_offset_sec))
        embeddings[str(index)] = list(map(float, segment.embeddings_float))

        # Create unique ID for each segment
        id_hash = hashlib.md5(f"{video_url}_{index}".encode()).hexdigest()

        document = {
            "id": id_hash,
            "fields": {
                "video_url": video_url,
                "title": title,
                "keywords": keywords,
                "video_summary": summary,
                "embedding_scope": segment.embedding_scope,
                "start_offset_sec": start_offsets,
                "end_offset_sec": end_offsets,
                "embeddings": embeddings,
            },
        }
        vespa_feed.append(document)

from vespa.io import VespaResponse

def callback(response: VespaResponse, id: str):
    if not response.is_successful():
        print(f"Failed to feed document {id} with status code {response.status_code}: {response.get_json()}")

# Feed data synchronously
app.feed_iterable(vespa_feed, schema="videos", callback=callback)
print("Data feeding complete!")

client = TwelveLabs(api_key=TL_API_KEY)
user_query = "Santa Claus on his sleigh"

# Generate embedding for the query
res = client.embed.create(
    model_name="Marengo-retrieval-2.7",
    text=user_query,
)

print("Created a text embedding")
print(f" Model: {res.model_name}")
if res.text_embedding is not None and res.text_embedding.segments is not None:
    q_embedding = res.text_embedding.segments[0].embeddings_float
    print(f" Embedding Dimension: {len(q_embedding)}")
    print(f" Sample 5 values from array: {q_embedding[:5]}")

with app.syncio(connections=1) as session:
    response: VespaQueryResponse = session.query(
        yql="select * from videos where userQuery() OR ({targetHits:100}nearestNeighbor(embeddings,q))",
        query=user_query,
        ranking="hybrid",
        hits=1,
        body={"input.query(q)": q_embedding},
    )
    assert response.is_successful()

# Print the top hit
for hit in response.hits:
    print(json.dumps(hit, indent=4))

# Get full response JSON
response.get_json()

def get_top_n_similarity_matches(data, N=5):
    """
    Extracts top N similarity scores and their corresponding offsets.

    Args:
    - data (dict): Input JSON-like structure containing similarities and offsets.
    - N (int): Number of top similarity scores to return.

    Returns:
    - pd.DataFrame: A DataFrame with top N similarity scores and offsets.
    """
    # Extract relevant fields
    similarities = data["fields"]["summaryfeatures"]["similarities"]["cells"]
    start_offset_sec = data["fields"]["start_offset_sec"]
    end_offset_sec = data["fields"]["end_offset_sec"]

    # Sort by similarity score
    sorted_similarities = sorted(similarities.items(), key=lambda x: x[1], reverse=True)

    # Prepare results for top N matches
    results = []
    for index_str, score in sorted_similarities[:N]:
        index = int(index_str)
        if index < len(start_offset_sec):
            results.append({
                "index": index,
                "similarity_score": score,
                "start_offset_sec": start_offset_sec[index],
                "end_offset_sec": end_offset_sec[index],
            })

    return pd.DataFrame(results)

# Get top 10 matches
df_result = get_top_n_similarity_matches(response.hits[0], N=10)
print(df_result)