マルチモーダルRAG：Twelve LabsとPineconeを使用したビデオとのチャット

# Install required libraries
!pip install twelvelabs pinecone-client

!pip install -q av
!pip install --upgrade -q accelerate bitsandbytes
!pip install transformers

from google.colab import userdata

TL_API_KEY=userdata.get('TL_API_KEY')
PINECONE_API_KEY=userdata.get('PINECONE_API_KEY')

from google.colab import drive
drive.mount('/content/drive')
base_folder_path = "/content/drive/MyDrive/TwelveLabs-Pinecone"
single_video = base_folder_path + "/ad_vids/Rare Beauty By Selena Gomez - Makeup Made To Feel Good In.mp4"
split_video_dir = base_folder_path + "/split_ad_videos"

# Configure Pinecone
from pinecone import Pinecone, ServerlessSpec

pc = Pinecone(api_key=PINECONE_API_KEY)
from twelvelabs import TwelveLabs
from twelvelabs.models.embed import EmbeddingsTask

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

# Define a callback function to monitor task progress
def on_task_update(task: EmbeddingsTask):
   print(f"  Status={task.status}")

def generate_embedding(video_file):
   # Create an embedding task
   task = twelvelabs_client.embed.task.create(
       engine_name="Marengo-retrieval-2.6",
       video_file=video_file
   )
   print(f"Created task: id={task.id} engine_name={task.engine_name} status={task.status}")


   # Wait for the task to complete
   status = task.wait_for_done(
       sleep_interval=2,
       callback=on_task_update
   )
   print(f"Embedding done: {status}")

   # Retrieve the task result
   task_result = twelvelabs_client.embed.task.retrieve(task.id)

   # Extract and return the embeddings
   embeddings = []
   for v in task_result.video_embeddings:
       embeddings.append({
           'embedding': v.embedding.float,
           'start_offset_sec': v.start_offset_sec,
           'end_offset_sec': v.end_offset_sec,
           'embedding_scope': v.embedding_scope
       })

   return embeddings, task_result


def ingest_data(video_file_path, index_name = "twelve-labs"):
   """
   Generate embeddings for video and store in Pinecone
   """

   #Strip the extension and the folders from the video_file_path
   video_name = os.path.splitext(os.path.basename(video_file_path))[0]
   print(video_name)
   # Connect to Pinecone index
   if index_name not in pc.list_indexes().names():
       pc.create_index(
           name=index_name,
           dimension=1024, # The dimensions of Twelve Lab's Embedding Model
           metric="cosine",
           spec=ServerlessSpec(
               cloud="aws",
               region="us-east-1"
           )

       )
   index = pc.Index(index_name)

   # Generate embeddings using Twelve Labs Embed API
   embeddings, task_result = generate_embedding(video_file_path)

   # Prepare vectors for upsert
   vectors_to_upsert = []
   for i, emb in enumerate(embeddings):
       vector_id = f"{video_name}_{i}"
       vectors_to_upsert.append((vector_id, emb['embedding'], {
           'video_file': video_name,
           'video_segment': i,
           'start_time': emb['start_offset_sec'],
           'end_time': emb['end_offset_sec'],
           'scope': emb['embedding_scope']
       }))

   # Upsert embeddings to Pinecone
   index.upsert(vectors=vectors_to_upsert)

   return f"Ingested {len(embeddings)} embeddings for {video_file_path}"

# Example usage
result = ingest_data(single_video)
print(result)

def get_text_embedding(text_query):
   # Twelve Labs Embed API supports text-to-embedding
   text_embedding = twelvelabs_client.embed.create(
     engine_name="Marengo-retrieval-2.6",
     text=text_query,
     text_truncate="start"
   )

   return text_embedding.text_embedding.float


def retrieve_similar_content(query, index_name="twelve-labs", top_k=5):
   """
   Retrieve similar content based on query embedding
   """
   # Generate query embedding
   query_embedding = get_text_embedding(query)

   # Connect to Pinecone index
   index = pc.Index(index_name)

   # Query Pinecone for similar vectors
   results = index.query(vector=query_embedding, top_k=top_k, include_metadata=True)

   return results

# Example usage
text_query = "Lipstick"
similar_content = retrieve_similar_content(text_query)

print(f"Query: '{text_query}'")
print(f"Top {len(similar_content['matches'])} similar content:")
for i, match in enumerate(similar_content['matches']):
   print(f"{i+1}. Score: {match['score']:.4f}")
   print(f"   Video File: {match['metadata']['video_file']}")
   print(f"   Video ID: {match['metadata']['video_segment']}")
   print(f"   Time range: {match['metadata']['start_time']} - {match['metadata']['end_time']} seconds")
   print(f"   Scope: {match['metadata']['scope']}")
   print()

import av

def split_video(input_path, output_dir, segment_duration=6):
   # Ensure output directory exists
   os.makedirs(output_dir, exist_ok=True)
   input_file_name = os.path.splitext(os.path.basename(input_path))[0]
   print(input_file_name)
   with av.open(input_path) as input_container:
       # Get video stream
       input_stream = input_container.streams.video[0]
       fps = input_stream.average_rate

       # Calculate how many frames are in each segment
       frames_per_segment = int(segment_duration * fps)

       segment_count = 0
       frame_count = 0
       output_container = None
       output_stream = None
       first_frame_timestamp = None

       for frame in input_container.decode(video=0):
           if frame_count % frames_per_segment == 0:
               # Close previous output container if it exists
               if output_container:
                   output_container.close()

               # Create a new output container

               output_path = os.path.join(output_dir, f'{input_file_name}_segment_{segment_count:03d}.mp4')
               segment_count += 1

               output_container = av.open(output_path, mode='w')
               output_stream = output_container.add_stream('h264', rate=fps)
               output_stream.width = frame.width
               output_stream.height = frame.height
               output_stream.pix_fmt = 'yuv420p'

               # Reset the first frame timestamp for the new segment
               first_frame_timestamp = frame.pts

           # Adjust the frame timestamp
           frame.pts -= first_frame_timestamp

           # Encode frame
           packet = output_stream.encode(frame)
           output_container.mux(packet)

           frame_count += 1

       # Flush the encoder
       packet = output_stream.encode(None)
       output_container.mux(packet)

       # Close the last output container
       if output_container:
           output_container.close()

split_video(input_path=single_video, output_dir=split_video_dir)

query = "What is this advertisement selling?"
similar_content = retrieve_similar_content(query)

engines = [
       {
           "name": "pegasus1.1",
           "options": ["visual", "conversation"]
       }
   ]

index_name = "ads_index"
indices_list = twelvelabs_client.index.list(name=index_name)

if len(indices_list) == 0:
   index = twelvelabs_client.index.create(
       name=index_name,
       engines=engines,

   )
   print(f"A new index has been created: id={index.id} name={index.name} engines={index.engines}")
else:
   index = indices_list[0]
   print(f"Index already exists: id={index.id} name={index.name} engines={index.engines}")

def upload_video_to_twelve_labs(video_path):
   task = twelvelabs_client.task.create(
       index_id=index.id,
       file = video_path
   )
   print(f"Task created: id={task.id} status={task.status}")

   task.wait_for_done(sleep_interval=5, callback=on_task_update)

   if task.status != "ready":
     raise RuntimeError(f"Indexing failed with status {task.status}")
   print(f"The unique identifier of your video is {task.video_id}.")
   return task.video_id

video_ids = {}
for split_video_filename in os.listdir(split_video_dir):
   split_video_path = os.path.join(split_video_dir, split_video_filename)
   print(split_video_path)
   split_video_name = split_video_filename.split('.')[0]
   print(split_video_name)
   video_id = upload_video_to_twelve_labs(split_video_path)
   video_ids[split_video_name] = video_id

print(video_ids)

# retrieve the correct video_id for the relevant video
video_segment = (int) (similar_content['matches'][0]['metadata']['video_segment'])
print(f"Retrieved video segment: {video_segment}")

base_filename = os.path.splitext(os.path.basename(single_video))[0]
video_key = f"{base_filename}_segment_{video_segment:03d}"

video_id = video_ids[video_key]

res = twelvelabs_client.generate.text(
 video_id=video_id,
 prompt=query
)
print(f"{res.data}")

import av
import numpy as np

def read_video_pyav(container, indices):
    '''
    Decode the video with PyAV decoder.

    Args:
        container (av.container.input.InputContainer): PyAV container.
        indices (List[int]): List of frame indices to decode.

    Returns:
        np.ndarray: np array of decoded frames of shape (num_frames, height, width, 3).
    '''
    frames = []
    container.seek(0)
    start_index = indices[0]
    end_index = indices[-1]
    for i, frame in enumerate(container.decode(video=0)):
        if i > end_index:
            break
        if i >= start_index and i in indices:
            frames.append(frame)

def get_total_frames(video_path):
    """
    Manually count the total number of frames in a video. Used in case uniformly sampling comes up as 0 frames.
    """
    container = av.open(video_path)
    video_stream = container.streams.video[0]
    total_frames = 0
    for frame in container.decode(video_stream):
        total_frames += 1

    return total_frames

def sample_video(video_path, num_samples=8):
    container = av.open(video_path)
    video_stream = container.streams.video[0]

    # sample uniformly num_samples frames from the video
    total_frames = container.streams.video[0].frames
    if total_frames == 0:
        total_frames = get_total_frames(video_path)
    indices = np.arange(0, total_frames, total_frames / num_samples).astype(int)
    sampled_frames = read_video_pyav(container, indices)

    return sampled_frames

def process_videos_in_folder(folder_path):
    sample_info = {}

    # Supported video file extensions
    video_extensions = ('.mp4', '.avi', '.mov', '.mkv')

    for filename in os.listdir(folder_path):
        simple_video_name = os.path.splitext(os.path.basename(filename))[0]
        if filename.lower().endswith(video_extensions):
            video_path = os.path.join(folder_path, filename)
            try:
                print("Sampling " + video_path)
                sampled_clip = sample_video(video_path)
                sample_info[simple_video_name] = {"sampled_video": sampled_clip, "video_path" : video_path}
            except Exception as e:
                print(f"Error processing {filename}: {str(e)}")

    return sample_info


sampled_video_info = process_videos_in_folder(split_video_dir)

from transformers import BitsAndBytesConfig, LlavaNextVideoForConditionalGeneration, LlavaNextVideoProcessor
import torch

quantization_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_compute_dtype=torch.float16
)

processor = LlavaNextVideoProcessor.from_pretrained("llava-hf/LLaVA-NeXT-Video-7B-hf")
model = LlavaNextVideoForConditionalGeneration.from_pretrained(
    "llava-hf/LLaVA-NeXT-Video-7B-hf",
    quantization_config=quantization_config,
    device_map='auto'
)

def video_segment_name_from_offset(video_path, start_time, segment_length = 6):
   segment_number = int (start_time // segment_length)
   simple_video_name = os.path.splitext(os.path.basename(video_path))[0]

   return f"{simple_video_name}_segment_{segment_number:03d}"

video_segment = similar_content['matches'][0]['metadata']['video_file']
print(video_segment)
video_offset = similar_content['matches'][0]['metadata']['start_time']
video_segment_name = video_segment_name_from_offset(video_segment, video_offset)
video_segment = sampled_video_info[video_segment_name]['sampled_video']

# Each "content" is a list of dicts and you can add image/video/text modalities
conversation = [
      {
          "role": "user",
          "content": [
              {"type": "text", "text": query},
              {"type": "video"},
              ],
      },
]

prompt = processor.apply_chat_template(conversation, add_generation_prompt=True)
prompt_len = len(prompt)

inputs = processor([prompt], videos=[video_segment], padding=True, return_tensors="pt").to(model.device)

generate_kwargs = {"max_new_tokens": 100, "do_sample": True, "top_p": 0.9}

output = model.generate(**inputs, **generate_kwargs)
generated_text = processor.batch_decode(output, skip_special_tokens=True)

print(generated_text[0])

ads_dir = os.path.join(base_folder_path,"ad_vids")
video_list = []
#Make sure that we' don't waste time re-embedding the original single video:
single_video_filename = os.path.splitext(os.path.basename(single_video))[0]

for filename in os.listdir(ads_dir):
  if filename.endswith(".mp4") and single_video_filename not in filename:
    video_list.append(ads_dir + "/" + filename)

print(video_list)
for video in video_list:
  ingest_data(video)

full_database_questions = ["Who is the actor in the Miss Dior video?", "What ad is Selena Gomez in?", "What is the ad for Rare Beauty about?", "Why should people buy the Rare Beauty product according to their ad?"]
question = full_database_questions[0]
similar_content_from_question = retrieve_similar_content(question)
video_name = similar_content_from_question['matches'][0]['metadata']['video_file']

for vid in os.listdir(ads_dir):
    vid_path = os.path.join(ads_dir, vid)
    vid_name = os.path.splitext(os.path.basename(vid_path))[0]
    print(vid_path)
    video_id = upload_video_to_twelve_labs(vid_path)
    video_ids[vid_name] = video_id

video_id = video_ids[video_name]

res = twelvelabs_client.generate.text(
  video_id=video_id,
  prompt=question
)
print(f"{res.data}")

sampled_database_video_info = process_videos_in_folder(ads_dir)
video_segment = sampled_database_video_info[video_name]['sampled_video']

conversation = [
      {
          "role": "user",
          "content": [
              {"type": "text", "text": question},
              {"type": "video"},
              ],
      },
]

prompt = processor.apply_chat_template(conversation, add_generation_prompt=True)
prompt_len = len(prompt)

inputs = processor([prompt], videos=[video_segment], padding=True, return_tensors="pt").to(model.device)

generate_kwargs = {"max_new_tokens": 100, "do_sample": True, "top_p": 0.9}

output = model.generate(**inputs, **generate_kwargs)
generated_text = processor.batch_decode(output, skip_special_tokens=True)

print(generated_text[0])