Add Marengo embedding predictor for zero-shot moment retrieval

README.md

@@ -94,6 +94,36 @@ Run `python api_example/amr_demo.py` to reproduce the AMR results.
 **Limitation**: The maximum video duration is **150s** due to the current benchmark datasets.
 For CPU users, set `feature_name='clip'` because CLIP+Slowfast or CLIP+Slowfast+PANNs features are very slow without GPUs.
 
+## Marengo embeddings (zero-shot, API-based)
+Lighthouse also ships an optional, zero-shot moment retrieval predictor backed by
+[TwelveLabs Marengo](https://twelvelabs.io) embeddings. It needs **no local checkpoint, feature
+files, or GPU**: the video is segmented and embedded server-side, the query is embedded into the
+same 512-dim space, and clips are ranked by cosine similarity. This is handy as a training-free
+baseline or for videos longer than the 150s benchmark limit.
+
+Install the extra and set your API key (create one in the
+[TwelveLabs dashboard](https://twelvelabs.io)):
+```
+pip install 'lighthouse[twelvelabs] @ git+https://github.com/line/lighthouse.git'
+export TWELVELABS_API_KEY=<your key>
+```
+```python
+from lighthouse.models import TwelveLabsPredictor
+
+model = TwelveLabsPredictor(clip_length=2.0)  # reads TWELVELABS_API_KEY from the env
+
+# a local file or a public video URL works
+video = model.encode_video('api_example/RoripwjYFp8_60.0_210.0.mp4')
+
+query = 'A woman wearing a glass is speaking in front of the camera'
+prediction = model.predict(query, video)
+print(prediction)
+"""
+{'pred_relevant_windows': [[start, end, score], ...]}  # sorted by score, same format as above
+"""
+```
+Run `python api_example/marengo_demo.py` to reproduce.
+
 ## Gradio demo
 Run `python gradio_demo/demo.py`. Upload the video and input text query, and click the blue button. For AMR demo, run `python gradio_demo/amr_demo.py`. 
 

api_example/marengo_demo.py

@@ -0,0 +1,33 @@
+"""
+Copyright $today.year LY Corporation
+
+LY Corporation licenses this file to you under the Apache License,
+version 2.0 (the "License"); you may not use this file except in compliance
+with the License. You may obtain a copy of the License at:
+
+  https://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+License for the specific language governing permissions and limitations
+under the License.
+"""
+import os
+
+from lighthouse.models import TwelveLabsPredictor
+from typing import Dict, List, Optional
+
+# Zero-shot moment retrieval with TwelveLabs Marengo embeddings.
+# No local checkpoint or feature files are needed; the video is embedded
+# server-side. Set TWELVELABS_API_KEY first (free key: https://twelvelabs.io).
+api_key: Optional[str] = os.environ.get('TWELVELABS_API_KEY')
+model: TwelveLabsPredictor = TwelveLabsPredictor(api_key=api_key, clip_length=2.0)
+
+# encode video clips (local file or a public URL)
+video: Dict[str, List[List[float]]] = model.encode_video('api_example/RoripwjYFp8_60.0_210.0.mp4')
+
+# moment retrieval
+query: str = 'A woman wearing a glass is speaking in front of the camera'
+prediction: Optional[Dict[str, List[List[float]]]] = model.predict(query, video)
+print(prediction)

lighthouse/marengo_predictor.py

@@ -0,0 +1,174 @@
+import os
+
+from typing import Dict, List, Optional
+
+"""
+Copyright $today.year LY Corporation
+
+LY Corporation licenses this file to you under the Apache License,
+version 2.0 (the "License"); you may not use this file except in compliance
+with the License. You may obtain a copy of the License at:
+
+  https://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+License for the specific language governing permissions and limitations
+under the License.
+"""
+
+
+class TwelveLabsPredictor:
+    """
+    Zero-shot video moment retrieval backed by TwelveLabs embeddings.
+
+    Unlike the other predictors in :mod:`lighthouse.models`, this one needs no
+    local checkpoint or feature files. It segments the video on the server side,
+    embeds each clip and the text query into the same 512-dim embedding space,
+    and ranks clips by cosine similarity. The output mirrors the rest of the
+    library:
+
+        {"pred_relevant_windows": [[start, end, score], ...]}
+
+    The embedding model is configurable via the ``model_name`` argument and
+    defaults to Marengo (``marengo3.0``), so the class can support future
+    TwelveLabs embedding models without code changes.
+
+    This is an opt-in addition; existing predictors and their behaviour are
+    unchanged. It depends on the official ``twelvelabs`` SDK, which is only
+    imported when this class is instantiated, so users who do not use it pay no
+    import cost.
+
+    A free API key with a generous free tier is available at https://twelvelabs.io.
+    """
+
+    DEFAULT_MODEL_NAME: str = 'marengo3.0'
+    EMBEDDING_DIM: int = 512
+
+    def __init__(
+        self,
+        api_key: Optional[str] = None,
+        model_name: str = DEFAULT_MODEL_NAME,
+        clip_length: float = 2.0,
+        moment_num: int = 10) -> None:
+
+        try:
+            from twelvelabs import TwelveLabs
+        except ImportError as e:
+            raise ImportError(
+                'TwelveLabsPredictor requires the twelvelabs SDK. '
+                "Install it with: pip install 'lighthouse[twelvelabs]' "
+                'or pip install twelvelabs.') from e
+
+        resolved_key = api_key or os.environ.get('TWELVELABS_API_KEY')
+        if not resolved_key:
+            raise ValueError(
+                'A TwelveLabs API key is required. Pass api_key=... or set the '
+                'TWELVELABS_API_KEY environment variable. Get a free key at '
+                'https://twelvelabs.io.')
+
+        self._client = TwelveLabs(api_key=resolved_key)
+        self.model_name: str = model_name
+        self._clip_length: float = clip_length
+        self._moment_num: int = moment_num
+
+    @staticmethod
+    def _cosine_similarity(
+        query: List[float],
+        clip: List[float]) -> float:
+        dot = sum(q * c for q, c in zip(query, clip))
+        query_norm = sum(q * q for q in query) ** 0.5
+        clip_norm = sum(c * c for c in clip) ** 0.5
+        if query_norm == 0.0 or clip_norm == 0.0:
+            return 0.0
+        return dot / (query_norm * clip_norm)
+
+    def encode_video(
+        self,
+        video_path: Optional[str] = None,
+        video_url: Optional[str] = None) -> Dict[str, List[List[float]]]:
+        """
+        Embed every clip of a video with TwelveLabs.
+
+        Provide either a local ``video_path`` or a publicly reachable
+        ``video_url``. Returns a dict with one ``segments`` entry, each being
+        ``[start_offset_sec, end_offset_sec, *embedding]``, ready to pass to
+        :meth:`predict`.
+        """
+        if (video_path is None) == (video_url is None):
+            raise ValueError('Pass exactly one of video_path or video_url.')
+
+        if video_path is not None:
+            with open(video_path, 'rb') as f:
+                task = self._client.embed.tasks.create(
+                    model_name=self.model_name,
+                    video_file=f,
+                    video_clip_length=self._clip_length,
+                    video_embedding_scope=['clip'])
+        else:
+            task = self._client.embed.tasks.create(
+                model_name=self.model_name,
+                video_url=video_url,
+                video_clip_length=self._clip_length,
+                video_embedding_scope=['clip'])
+
+        if task.id is None:
+            raise RuntimeError('The embedding task creation returned no task id.')
+        task_id: str = task.id
+
+        # wait_for_done lives on the SDK's runtime wrapper class, which mypy
+        # cannot see through the statically-typed base client.
+        self._client.embed.tasks.wait_for_done(task_id=task_id)  # type: ignore[attr-defined]
+        result = self._client.embed.tasks.retrieve(task_id=task_id)
+
+        if result.video_embedding is None or result.video_embedding.segments is None:
+            raise RuntimeError(f'The embedding task {task_id} returned no segments.')
+
+        segments: List[List[float]] = []
+        for segment in result.video_embedding.segments:
+            if segment.float_ is None:
+                continue
+            start = float(segment.start_offset_sec or 0.0)
+            end = float(segment.end_offset_sec or 0.0)
+            segments.append([start, end] + list(segment.float_))
+        return {'segments': segments}
+
+    def _encode_text(
+        self,
+        query: str) -> List[float]:
+        response = self._client.embed.create(model_name=self.model_name, text=query)
+        if response.text_embedding is None or not response.text_embedding.segments:
+            raise RuntimeError('TwelveLabs returned no text embedding for the query.')
+        embedding = response.text_embedding.segments[0].float_
+        if embedding is None:
+            raise RuntimeError('TwelveLabs returned an empty text embedding for the query.')
+        return list(embedding)
+
+    def predict(
+        self,
+        query: str,
+        inputs: Dict[str, List[List[float]]]) -> Optional[Dict[str, List[List[float]]]]:
+        """
+        Rank the encoded video clips by similarity to ``query``.
+
+        Returns ``{"pred_relevant_windows": [[start, end, score], ...]}`` sorted
+        by descending score, capped at ``moment_num`` windows, matching the
+        prediction format of the other Lighthouse predictors.
+        """
+        segments = inputs.get('segments')
+        if not segments:
+            print('Error: No encoded clips found. Did you forget to call encode_video()?')
+            return None
+
+        query_embedding = self._encode_text(query)
+
+        ranked: List[List[float]] = []
+        for segment in segments:
+            start, end = segment[0], segment[1]
+            clip_embedding = segment[2:]
+            score = self._cosine_similarity(query_embedding, clip_embedding)
+            ranked.append([float(f'{start:.4f}'), float(f'{end:.4f}'), float(f'{score:.4f}')])
+
+        ranked.sort(key=lambda window: window[2], reverse=True)
+        return {'pred_relevant_windows': ranked[:self._moment_num]}

lighthouse/models.py

@@ -13,6 +13,7 @@
 from lighthouse.feature_extractor.vision_encoder import VisionEncoder
 from lighthouse.feature_extractor.text_encoder import TextEncoder
 from lighthouse.feature_extractor.audio_encoder import AudioEncoder
+from lighthouse.marengo_predictor import TwelveLabsPredictor as TwelveLabsPredictor
 
 from typing import Optional, Union, Mapping, Any, Dict, List, Tuple
 

mypy.ini

@@ -33,3 +33,6 @@ ignore_missing_imports = True
 
 [mypy-msclap.*]
 ignore_missing_imports = True
+
+[mypy-twelvelabs.*]
+ignore_missing_imports = True

setup.py

@@ -6,5 +6,6 @@
     install_requires=['easydict', 'pandas', 'tqdm', 'pyyaml', 'scikit-learn', 'ffmpeg-python',
                       'ftfy', 'regex', 'einops', 'fvcore', 'gradio', 'torchlibrosa', 'librosa',
                       'msclap', 'transformers<=4.51.3', 'numpy<=1.23.5', 'clip@git+https://github.com/openai/CLIP.git'],
+    extras_require={'twelvelabs': ['twelvelabs>=1.2.8']},
     packages=find_packages(exclude=['training']),
 )

tests/test_marengo.py

@@ -0,0 +1,48 @@
+import os
+import pytest
+
+from lighthouse.models import TwelveLabsPredictor
+
+
+MOMENT_NUM = 10
+EMBEDDING_DIM = 512
+
+
+def test_predict_ranks_clips_by_cosine_similarity():
+    """No-network unit test: predict() must rank clips by query similarity."""
+    predictor = TwelveLabsPredictor.__new__(TwelveLabsPredictor)
+    predictor._moment_num = MOMENT_NUM
+
+    # clip 0 points the same way as the query, clip 1 is orthogonal.
+    inputs = {
+        'segments': [
+            [0.0, 2.0, 1.0, 0.0],
+            [2.0, 4.0, 0.0, 1.0],
+        ],
+    }
+    predictor._encode_text = lambda query: [1.0, 0.0]  # type: ignore[method-assign]
+
+    prediction = predictor.predict('a query', inputs)
+    windows = prediction['pred_relevant_windows']
+
+    assert len(windows) == 2
+    assert windows[0][:2] == [0.0, 2.0], 'the aligned clip should rank first'
+    assert windows[0][2] == 1.0
+    assert windows[1][2] == 0.0
+    assert windows[0][2] >= windows[1][2], 'windows must be sorted by descending score'
+
+
+def test_predict_without_segments_returns_none():
+    predictor = TwelveLabsPredictor.__new__(TwelveLabsPredictor)
+    predictor._moment_num = MOMENT_NUM
+    assert predictor.predict('a query', {'segments': []}) is None
+
+
+@pytest.mark.skipif(
+    'TWELVELABS_API_KEY' not in os.environ,
+    reason='TWELVELABS_API_KEY not set; skipping live Marengo API test.')
+def test_marengo_text_embedding_is_512_dim():
+    """Live smoke test: a Marengo text embedding is a 512-dim vector."""
+    predictor = TwelveLabsPredictor(clip_length=2.0)
+    embedding = predictor._encode_text('a person walking a dog')
+    assert len(embedding) == EMBEDDING_DIM