wehub-resource-sync
430 lines
16 KiB
Python
430 lines
16 KiB
Python
# Copyright (c) Facebook, Inc. and its affiliates.
|
|
#
|
|
# This source code is licensed under the MIT license found in the
|
|
# LICENSE file in the root directory of this source tree.
|
|
import os
|
|
import numpy as np
|
|
import pickle
|
|
import time
|
|
|
|
try:
|
|
import faiss
|
|
except ImportError:
|
|
pass
|
|
|
|
from collections import defaultdict
|
|
|
|
from ..utils import get_local_rank, print_on_rank0
|
|
|
|
|
|
class VectorRetriever(object):
|
|
"""
|
|
How2 Video Retriver.
|
|
Reference usage of FAISS:
|
|
https://github.com/fairinternal/fairseq-py/blob/paraphrase_pretraining/fairseq/data/multilingual_faiss_dataset.py
|
|
"""
|
|
|
|
def __init__(self, hidden_size, cent, db_type, examples_per_cent_to_train):
|
|
if db_type == "flatl2":
|
|
quantizer = faiss.IndexFlatL2(hidden_size) # the other index
|
|
self.db = faiss.IndexIVFFlat(
|
|
quantizer, hidden_size, cent, faiss.METRIC_L2)
|
|
elif db_type == "pq":
|
|
self.db = faiss.index_factory(
|
|
hidden_size, f"IVF{cent}_HNSW32,PQ32"
|
|
)
|
|
else:
|
|
raise ValueError("unknown type of db", db_type)
|
|
self.train_thres = cent * examples_per_cent_to_train
|
|
self.train_cache = []
|
|
self.train_len = 0
|
|
self.videoid_to_vectoridx = {}
|
|
self.vectoridx_to_videoid = None
|
|
self.make_direct_maps_done = False
|
|
|
|
def make_direct_maps(self):
|
|
faiss.downcast_index(self.db).make_direct_map()
|
|
|
|
def __len__(self):
|
|
return self.db.ntotal
|
|
|
|
def save(self, out_dir):
|
|
faiss.write_index(
|
|
self.db,
|
|
os.path.join(out_dir, "faiss_idx")
|
|
)
|
|
with open(
|
|
os.path.join(
|
|
out_dir, "videoid_to_vectoridx.pkl"),
|
|
"wb") as fw:
|
|
pickle.dump(
|
|
self.videoid_to_vectoridx, fw,
|
|
protocol=pickle.HIGHEST_PROTOCOL
|
|
)
|
|
|
|
def load(self, out_dir):
|
|
fn = os.path.join(out_dir, "faiss_idx")
|
|
self.db = faiss.read_index(fn)
|
|
with open(
|
|
os.path.join(out_dir, "videoid_to_vectoridx.pkl"), "rb") as fr:
|
|
self.videoid_to_vectoridx = pickle.load(fr)
|
|
|
|
def add(self, hidden_states, video_ids, last=False):
|
|
assert len(hidden_states) == len(video_ids), "{}, {}".format(
|
|
str(len(hidden_states)), str(len(video_ids)))
|
|
assert len(hidden_states.shape) == 2
|
|
assert hidden_states.dtype == np.float32
|
|
|
|
valid_idx = []
|
|
for idx, video_id in enumerate(video_ids):
|
|
if video_id not in self.videoid_to_vectoridx:
|
|
valid_idx.append(idx)
|
|
self.videoid_to_vectoridx[video_id] = \
|
|
len(self.videoid_to_vectoridx)
|
|
|
|
hidden_states = hidden_states[valid_idx]
|
|
if not self.db.is_trained:
|
|
self.train_cache.append(hidden_states)
|
|
self.train_len += hidden_states.shape[0]
|
|
if self.train_len < self.train_thres:
|
|
return
|
|
self.finalize_training()
|
|
else:
|
|
self.db.add(hidden_states)
|
|
|
|
def finalize_training(self):
|
|
hidden_states = np.concatenate(self.train_cache, axis=0)
|
|
del self.train_cache
|
|
local_rank = get_local_rank()
|
|
if local_rank == 0:
|
|
start = time.time()
|
|
print("training db on", self.train_thres, "/", self.train_len)
|
|
self.db.train(hidden_states[:self.train_thres])
|
|
if local_rank == 0:
|
|
print("training db for", time.time() - start)
|
|
self.db.add(hidden_states)
|
|
|
|
def search(
|
|
self,
|
|
query_hidden_states,
|
|
orig_dist,
|
|
):
|
|
if len(self.videoid_to_vectoridx) != self.db.ntotal:
|
|
raise ValueError(
|
|
"cannot search: size mismatch in-between index and db",
|
|
len(self.videoid_to_vectoridx),
|
|
self.db.ntotal
|
|
)
|
|
|
|
if self.vectoridx_to_videoid is None:
|
|
self.vectoridx_to_videoid = {
|
|
self.videoid_to_vectoridx[videoid]: videoid
|
|
for videoid in self.videoid_to_vectoridx
|
|
}
|
|
assert len(self.vectoridx_to_videoid) \
|
|
== len(self.videoid_to_vectoridx)
|
|
|
|
# MultilingualFaissDataset uses the following; not sure the purpose.
|
|
# faiss.ParameterSpace().set_index_parameter(self.db, "nprobe", 10)
|
|
queried_dist, index = self.db.search(query_hidden_states, 1)
|
|
queried_dist, index = queried_dist[:, 0], index[:, 0]
|
|
|
|
outputs = np.array(
|
|
[self.vectoridx_to_videoid[_index]
|
|
if _index != -1 else (-1, -1, -1) for _index in index],
|
|
dtype=np.int32)
|
|
outputs[queried_dist <= orig_dist] = -1
|
|
return outputs
|
|
|
|
def search_by_video_ids(
|
|
self,
|
|
video_ids,
|
|
retri_factor
|
|
):
|
|
if len(self.videoid_to_vectoridx) != self.db.ntotal:
|
|
raise ValueError(
|
|
len(self.videoid_to_vectoridx),
|
|
self.db.ntotal
|
|
)
|
|
|
|
if not self.make_direct_maps_done:
|
|
self.make_direct_maps()
|
|
|
|
if self.vectoridx_to_videoid is None:
|
|
self.vectoridx_to_videoid = {
|
|
self.videoid_to_vectoridx[videoid]: videoid
|
|
for videoid in self.videoid_to_vectoridx
|
|
}
|
|
assert len(self.vectoridx_to_videoid) \
|
|
== len(self.videoid_to_vectoridx)
|
|
|
|
query_hidden_states = []
|
|
vector_ids = []
|
|
for video_id in video_ids:
|
|
vector_id = self.videoid_to_vectoridx[video_id]
|
|
vector_ids.append(vector_id)
|
|
query_hidden_state = self.db.reconstruct(vector_id)
|
|
query_hidden_states.append(query_hidden_state)
|
|
query_hidden_states = np.stack(query_hidden_states)
|
|
|
|
# MultilingualFaissDataset uses the following; not sure the reason.
|
|
# faiss.ParameterSpace().set_index_parameter(self.db, "nprobe", 10)
|
|
_, index = self.db.search(query_hidden_states, retri_factor)
|
|
outputs = []
|
|
for sample_idx, sample in enumerate(index):
|
|
# the first video_id is always the video itself.
|
|
cands = [video_ids[sample_idx]]
|
|
for vector_idx in sample:
|
|
if vector_idx >= 0 \
|
|
and vector_ids[sample_idx] != vector_idx:
|
|
cands.append(
|
|
self.vectoridx_to_videoid[vector_idx]
|
|
)
|
|
outputs.append(cands)
|
|
return outputs
|
|
|
|
|
|
class VectorRetrieverDM(VectorRetriever):
|
|
"""
|
|
with direct map.
|
|
How2 Video Retriver.
|
|
Reference usage of FAISS:
|
|
https://github.com/fairinternal/fairseq-py/blob/paraphrase_pretraining/fairseq/data/multilingual_faiss_dataset.py
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
hidden_size,
|
|
cent,
|
|
db_type,
|
|
examples_per_cent_to_train
|
|
):
|
|
super().__init__(
|
|
hidden_size, cent, db_type, examples_per_cent_to_train)
|
|
self.make_direct_maps_done = False
|
|
|
|
def make_direct_maps(self):
|
|
faiss.downcast_index(self.db).make_direct_map()
|
|
self.make_direct_maps_done = True
|
|
|
|
def search(
|
|
self,
|
|
query_hidden_states,
|
|
orig_dist,
|
|
):
|
|
if len(self.videoid_to_vectoridx) != self.db.ntotal:
|
|
raise ValueError(
|
|
len(self.videoid_to_vectoridx),
|
|
self.db.ntotal
|
|
)
|
|
|
|
if not self.make_direct_maps_done:
|
|
self.make_direct_maps()
|
|
if self.vectoridx_to_videoid is None:
|
|
self.vectoridx_to_videoid = {
|
|
self.videoid_to_vectoridx[videoid]: videoid
|
|
for videoid in self.videoid_to_vectoridx
|
|
}
|
|
assert len(self.vectoridx_to_videoid) \
|
|
== len(self.videoid_to_vectoridx)
|
|
|
|
# MultilingualFaissDataset uses the following; not sure the reason.
|
|
# faiss.ParameterSpace().set_index_parameter(self.db, "nprobe", 10)
|
|
queried_dist, index = self.db.search(query_hidden_states, 1)
|
|
outputs = []
|
|
for sample_idx, sample in enumerate(index):
|
|
# and queried_dist[sample_idx] < thres \
|
|
if sample >= 0 \
|
|
and queried_dist[sample_idx] < orig_dist[sample_idx]:
|
|
outputs.append(self.vectoridx_to_videoid[sample])
|
|
else:
|
|
outputs.append(None)
|
|
return outputs
|
|
|
|
def search_by_video_ids(
|
|
self,
|
|
video_ids,
|
|
retri_factor=8
|
|
):
|
|
if len(self.videoid_to_vectoridx) != self.db.ntotal:
|
|
raise ValueError(
|
|
len(self.videoid_to_vectoridx),
|
|
self.db.ntotal
|
|
)
|
|
|
|
if not self.make_direct_maps_done:
|
|
self.make_direct_maps()
|
|
if self.vectoridx_to_videoid is None:
|
|
self.vectoridx_to_videoid = {
|
|
self.videoid_to_vectoridx[videoid]: videoid
|
|
for videoid in self.videoid_to_vectoridx
|
|
}
|
|
assert len(self.vectoridx_to_videoid) \
|
|
== len(self.videoid_to_vectoridx)
|
|
|
|
query_hidden_states = []
|
|
vector_ids = []
|
|
for video_id in video_ids:
|
|
vector_id = self.videoid_to_vectoridx[video_id]
|
|
vector_ids.append(vector_id)
|
|
query_hidden_state = self.db.reconstruct(vector_id)
|
|
query_hidden_states.append(query_hidden_state)
|
|
query_hidden_states = np.stack(query_hidden_states)
|
|
|
|
# MultilingualFaissDataset uses the following; not sure the reason.
|
|
# faiss.ParameterSpace().set_index_parameter(self.db, "nprobe", 10)
|
|
_, index = self.db.search(query_hidden_states, retri_factor)
|
|
outputs = []
|
|
for sample_idx, sample in enumerate(index):
|
|
# the first video_id is always the video itself.
|
|
cands = [video_ids[sample_idx]]
|
|
for vector_idx in sample:
|
|
if vector_idx >= 0 \
|
|
and vector_ids[sample_idx] != vector_idx:
|
|
cands.append(
|
|
self.vectoridx_to_videoid[vector_idx]
|
|
)
|
|
outputs.append(cands)
|
|
return outputs
|
|
|
|
|
|
class MMVectorRetriever(VectorRetrieverDM):
|
|
"""
|
|
multimodal vector retriver:
|
|
text retrieve video or video retrieve text.
|
|
"""
|
|
|
|
def __init__(self, hidden_size, cent, db_type, examples_per_cent_to_train):
|
|
super().__init__(
|
|
hidden_size, cent, db_type, examples_per_cent_to_train)
|
|
video_db = self.db
|
|
super().__init__(
|
|
hidden_size, cent, db_type, examples_per_cent_to_train)
|
|
text_db = self.db
|
|
self.db = {"video": video_db, "text": text_db}
|
|
self.video_to_videoid = defaultdict(list)
|
|
|
|
def __len__(self):
|
|
assert self.db["video"].ntotal == self.db["text"].ntotal
|
|
return self.db["video"].ntotal
|
|
|
|
def make_direct_maps(self):
|
|
faiss.downcast_index(self.db["video"]).make_direct_map()
|
|
faiss.downcast_index(self.db["text"]).make_direct_map()
|
|
|
|
def save(self, out_dir):
|
|
faiss.write_index(
|
|
self.db["video"],
|
|
os.path.join(out_dir, "video_faiss_idx")
|
|
)
|
|
faiss.write_index(
|
|
self.db["text"],
|
|
os.path.join(out_dir, "text_faiss_idx")
|
|
)
|
|
|
|
with open(
|
|
os.path.join(
|
|
out_dir, "videoid_to_vectoridx.pkl"),
|
|
"wb") as fw:
|
|
pickle.dump(
|
|
self.videoid_to_vectoridx, fw,
|
|
protocol=pickle.HIGHEST_PROTOCOL
|
|
)
|
|
|
|
def load(self, out_dir):
|
|
fn = os.path.join(out_dir, "video_faiss_idx")
|
|
video_db = faiss.read_index(fn)
|
|
fn = os.path.join(out_dir, "text_faiss_idx")
|
|
text_db = faiss.read_index(fn)
|
|
self.db = {"video": video_db, "text": text_db}
|
|
with open(
|
|
os.path.join(out_dir, "videoid_to_vectoridx.pkl"), "rb") as fr:
|
|
self.videoid_to_vectoridx = pickle.load(fr)
|
|
self.video_to_videoid = defaultdict(list)
|
|
|
|
def add(self, hidden_states, video_ids):
|
|
"""hidden_states is a pair `(video, text)`"""
|
|
assert len(hidden_states) == len(video_ids), "{}, {}".format(
|
|
str(len(hidden_states)), str(len(video_ids)))
|
|
assert len(hidden_states.shape) == 3
|
|
assert len(self.video_to_videoid) == 0
|
|
|
|
valid_idx = []
|
|
for idx, video_id in enumerate(video_ids):
|
|
if video_id not in self.videoid_to_vectoridx:
|
|
valid_idx.append(idx)
|
|
self.videoid_to_vectoridx[video_id] = \
|
|
len(self.videoid_to_vectoridx)
|
|
|
|
batch_size = hidden_states.shape[0]
|
|
hidden_states = hidden_states[valid_idx]
|
|
|
|
hidden_states = np.transpose(hidden_states, (1, 0, 2)).copy()
|
|
if not self.db["video"].is_trained:
|
|
self.train_cache.append(hidden_states)
|
|
train_len = batch_size * len(self.train_cache)
|
|
if train_len < self.train_thres:
|
|
return
|
|
|
|
hidden_states = np.concatenate(self.train_cache, axis=1)
|
|
del self.train_cache
|
|
self.db["video"].train(hidden_states[0, :self.train_thres])
|
|
self.db["text"].train(hidden_states[1, :self.train_thres])
|
|
self.db["video"].add(hidden_states[0])
|
|
self.db["text"].add(hidden_states[1])
|
|
|
|
def get_clips_by_video_id(self, video_id):
|
|
if not self.video_to_videoid:
|
|
for video_id, video_clip, text_clip in self.videoid_to_vectoridx:
|
|
self.video_to_videoid[video_id].append(
|
|
(video_id, video_clip, text_clip))
|
|
return self.video_to_videoid[video_id]
|
|
|
|
def search(
|
|
self,
|
|
video_ids,
|
|
target_modality,
|
|
retri_factor=8
|
|
):
|
|
if len(self.videoid_to_vectoridx) != len(self):
|
|
raise ValueError(
|
|
len(self.videoid_to_vectoridx),
|
|
len(self)
|
|
)
|
|
|
|
if not self.make_direct_maps_done:
|
|
self.make_direct_maps()
|
|
if self.vectoridx_to_videoid is None:
|
|
self.vectoridx_to_videoid = {
|
|
self.videoid_to_vectoridx[videoid]: videoid
|
|
for videoid in self.videoid_to_vectoridx
|
|
}
|
|
assert len(self.vectoridx_to_videoid) \
|
|
== len(self.videoid_to_vectoridx)
|
|
|
|
src_modality = "text" if target_modality == "video" else "video"
|
|
|
|
query_hidden_states = []
|
|
vector_ids = []
|
|
for video_id in video_ids:
|
|
vector_id = self.videoid_to_vectoridx[video_id]
|
|
vector_ids.append(vector_id)
|
|
query_hidden_state = self.db[src_modality].reconstruct(vector_id)
|
|
query_hidden_states.append(query_hidden_state)
|
|
query_hidden_states = np.stack(query_hidden_states)
|
|
|
|
# MultilingualFaissDataset uses the following; not sure the reason.
|
|
# faiss.ParameterSpace().set_index_parameter(self.db, "nprobe", 10)
|
|
_, index = self.db[target_modality].search(
|
|
query_hidden_states, retri_factor)
|
|
outputs = []
|
|
for sample_idx, sample in enumerate(index):
|
|
cands = []
|
|
for vector_idx in sample:
|
|
if vector_idx >= 0:
|
|
cands.append(
|
|
self.vectoridx_to_videoid[vector_idx]
|
|
)
|
|
outputs.append(cands)
|
|
return outputs
|