changing to x-clip

src/embedder.py

@@ -1,28 +1,26 @@
 import torch
 import numpy as np
-from transformers import AutoProcessor, AutoModel
-import decord
+from transformers import XCLIPProcessor, XCLIPModel
+from decord import VideoReader, cpu
 
-class InternVLEmbedder:
-    def __init__(self):
-        self.device = "cuda" if torch.cuda.is_available() else "cpu"
-        self.model = AutoModel.from_pretrained("OpenGVLab/InternVL2_5-1B-MPO", trust_remote_code=True).to(self.device)
-        self.processor = AutoProcessor.from_pretrained("OpenGVLab/InternVL2_5-1B-MPO", trust_remote_code=True)
+class XCLIPEmbedder:
+    def __init__(self, model_name="microsoft/xclip-base-patch14"):
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.model = XCLIPModel.from_pretrained(model_name).to(self.device)
+        self.processor = XCLIPProcessor.from_pretrained(model_name)
 
     def embed_video(self, video_path):
-        vr = decord.VideoReader(video_path)
-        frames = np.stack([vr[i].asnumpy() for i in np.linspace(0, len(vr)-1, 8).astype(int)])
-        tensor = torch.tensor(frames).permute(0, 3, 1, 2).unsqueeze(0).to(self.device)
+        vr = VideoReader(video_path, ctx=cpu(0))
+        frame_indices = np.linspace(0, len(vr) - 1, num=8, dtype=int)
+        video_frames = vr.get_batch(frame_indices).asnumpy()
 
+        inputs = self.processor(videos=list(video_frames), return_tensors="pt", padding=True).to(self.device)
         with torch.no_grad():
-            video_vector = self.model.get_video_features(tensor).squeeze(0).cpu().numpy()
-
-        return video_vector / np.linalg.norm(video_vector)
+            video_features = self.model.get_video_features(**inputs).squeeze(0).cpu().numpy()
+        return video_features / np.linalg.norm(video_features)
 
     def embed_text(self, text):
-        inputs = self.processor(text=[text], return_tensors="pt").to(self.device)
-
+        inputs = self.processor(text=[text], return_tensors="pt", padding=True).to(self.device)
         with torch.no_grad():
-            text_vector = self.model.get_text_features(**inputs).squeeze(0).cpu().numpy()
-
-        return text_vector / np.linalg.norm(text_vector)
+            text_features = self.model.get_text_features(**inputs).squeeze(0).cpu().numpy()
+        return text_features / np.linalg.norm(text_features)

src/pipeline.py

@@ -2,7 +2,7 @@ from src.segmenter import detect_event_segments
 from src.transcriber import transcribe_video
 from src.event_card import parse_game_card
 from src.labeler import TogetherLLMLabeler
-from src.embedder import InternVLEmbedder
+from src.embedder import XCLIPEmbedder
 from src.pinecone_store import PineconeStore
 from src.utils import (
     extract_key_frames, save_frames_locally,
@@ -11,7 +11,7 @@ from src.utils import (
 )
 
 labeler = TogetherLLMLabeler()
-embedder = InternVLEmbedder()
+embedder = XCLIPEmbedder()
 pinecone = PineconeStore()
 
 def run_pipeline(video_path, game_card_str):
@@ -40,7 +40,6 @@ def run_pipeline(video_path, game_card_str):
         clip_path = clip_video_segment(video_path, event['start_sec'], event['end_sec'], event_id)
 
         video_vector = embedder.embed_video(clip_path)
-        text_vector = embedder.embed_text(label)
 
         metadata = {
             "start_sec": event['start_sec'],
@@ -49,7 +48,6 @@ def run_pipeline(video_path, game_card_str):
         }
 
         pinecone.upsert(f"{event_id}-video", video_vector, metadata)
-        pinecone.upsert(f"{event_id}-text", text_vector, metadata)
 
         results.append(metadata)
 
@@ -57,7 +55,7 @@ def run_pipeline(video_path, game_card_str):
 
 def search_highlights(query, top_k=5):
     query_vector = embedder.embed_text(query)
-    results = pinecone.query(query_vector, filter_key="text", top_k=top_k)
+    results = pinecone.query(query_vector, filter_key="video", top_k=top_k)
     return [
         f"{r['label']} ({r['start_sec']}s - {r['end_sec']}s)" for r in results
     ]