mirror of
https://github.com/meta-llama/llama-stack.git
synced 2025-08-01 16:24:44 +00:00
fix faiss serialize and serialize of index
This commit is contained in:
parent
ff99025875
commit
fe40cdf28e
1 changed files with 8 additions and 6 deletions
|
@ -5,6 +5,7 @@
|
||||||
# the root directory of this source tree.
|
# the root directory of this source tree.
|
||||||
|
|
||||||
import base64
|
import base64
|
||||||
|
import io
|
||||||
import json
|
import json
|
||||||
import logging
|
import logging
|
||||||
|
|
||||||
|
@ -67,19 +68,20 @@ class FaissIndex(EmbeddingIndex):
|
||||||
for k, v in data["chunk_by_index"].items()
|
for k, v in data["chunk_by_index"].items()
|
||||||
}
|
}
|
||||||
|
|
||||||
index_bytes = base64.b64decode(data["faiss_index"])
|
buffer = io.BytesIO(base64.b64decode(data["faiss_index"]))
|
||||||
self.index = faiss.deserialize_index(index_bytes)
|
self.index = faiss.deserialize_index(np.loadtxt(buffer, dtype=np.uint8))
|
||||||
|
|
||||||
async def _save_index(self):
|
async def _save_index(self):
|
||||||
if not self.kvstore or not self.bank_id:
|
if not self.kvstore or not self.bank_id:
|
||||||
return
|
return
|
||||||
|
|
||||||
index_bytes = faiss.serialize_index(self.index)
|
np_index = faiss.serialize_index(self.index)
|
||||||
|
buffer = io.BytesIO()
|
||||||
|
np.savetxt(buffer, np_index)
|
||||||
data = {
|
data = {
|
||||||
"id_by_index": self.id_by_index,
|
"id_by_index": self.id_by_index,
|
||||||
"chunk_by_index": {k: v.json() for k, v in self.chunk_by_index.items()},
|
"chunk_by_index": {k: v.json() for k, v in self.chunk_by_index.items()},
|
||||||
"faiss_index": base64.b64encode(index_bytes).decode(),
|
"faiss_index": base64.b64encode(buffer.getvalue()).decode("utf-8"),
|
||||||
}
|
}
|
||||||
|
|
||||||
index_key = f"faiss_index:v1::{self.bank_id}"
|
index_key = f"faiss_index:v1::{self.bank_id}"
|
||||||
|
@ -188,7 +190,7 @@ class FaissMemoryImpl(Memory, MemoryBanksProtocolPrivate):
|
||||||
) -> None:
|
) -> None:
|
||||||
index = self.cache.get(bank_id)
|
index = self.cache.get(bank_id)
|
||||||
if index is None:
|
if index is None:
|
||||||
raise ValueError(f"Bank {bank_id} not found")
|
raise ValueError(f"Bank {bank_id} not found. found: {self.cache.keys()}")
|
||||||
|
|
||||||
await index.insert_documents(documents)
|
await index.insert_documents(documents)
|
||||||
|
|
||||||
|
|
Loading…
Add table
Add a link
Reference in a new issue