21 lines
766 B
Python
21 lines
766 B
Python
import os
|
|
import faiss
|
|
|
|
class UserIndexManager:
|
|
def __init__(self, base_dir="index_data/"):
|
|
self.base_dir = base_dir
|
|
self.index_map = {} # user_id -> faiss index
|
|
self._load_all_indexes()
|
|
|
|
def _load_all_indexes(self):
|
|
for fname in os.listdir(self.base_dir):
|
|
if fname.endswith(".index"):
|
|
user_id = fname.replace(".index", "")
|
|
path = os.path.join(self.base_dir, fname)
|
|
self.index_map[user_id] = faiss.read_index(path)
|
|
print(f"[INIT] Loaded index for user {user_id}")
|
|
|
|
def get_index(self, user_id):
|
|
if user_id not in self.index_map:
|
|
raise ValueError(f"Index for user {user_id} not loaded.")
|
|
return self.index_map[user_id] |