forked from phoenix-oss/llama-stack-mirror
120 lines
3.4 KiB
Python
120 lines
3.4 KiB
Python
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
|
# All rights reserved.
|
|
#
|
|
# This source code is licensed under the terms described in the LICENSE file in
|
|
# the root directory of this source tree.
|
|
|
|
from abc import ABC, abstractmethod
|
|
from dataclasses import dataclass
|
|
from typing import Any, Dict, List, Optional
|
|
|
|
import httpx
|
|
import numpy as np
|
|
from numpy.typing import NDArray
|
|
|
|
from llama_models.llama3.api.datatypes import * # noqa: F403
|
|
from llama_models.llama3.api.tokenizer import Tokenizer
|
|
|
|
from llama_toolchain.memory.api import * # noqa: F403
|
|
|
|
|
|
ALL_MINILM_L6_V2_DIMENSION = 384
|
|
|
|
EMBEDDING_MODEL = None
|
|
|
|
|
|
def get_embedding_model() -> "SentenceTransformer":
|
|
global EMBEDDING_MODEL
|
|
|
|
if EMBEDDING_MODEL is None:
|
|
print("Loading sentence transformer")
|
|
|
|
from sentence_transformers import SentenceTransformer
|
|
|
|
EMBEDDING_MODEL = SentenceTransformer("all-MiniLM-L6-v2")
|
|
|
|
return EMBEDDING_MODEL
|
|
|
|
|
|
async def content_from_doc(doc: MemoryBankDocument) -> str:
|
|
if isinstance(doc.content, URL):
|
|
async with httpx.AsyncClient() as client:
|
|
r = await client.get(doc.content.uri)
|
|
return r.text
|
|
|
|
return interleaved_text_media_as_str(doc.content)
|
|
|
|
|
|
def make_overlapped_chunks(
|
|
document_id: str, text: str, window_len: int, overlap_len: int
|
|
) -> List[Chunk]:
|
|
tokenizer = Tokenizer.get_instance()
|
|
tokens = tokenizer.encode(text, bos=False, eos=False)
|
|
|
|
chunks = []
|
|
for i in range(0, len(tokens), window_len - overlap_len):
|
|
toks = tokens[i : i + window_len]
|
|
chunk = tokenizer.decode(toks)
|
|
chunks.append(
|
|
Chunk(content=chunk, token_count=len(toks), document_id=document_id)
|
|
)
|
|
|
|
return chunks
|
|
|
|
|
|
class EmbeddingIndex(ABC):
|
|
@abstractmethod
|
|
async def add_chunks(self, chunks: List[Chunk], embeddings: NDArray):
|
|
raise NotImplementedError()
|
|
|
|
@abstractmethod
|
|
async def query(self, embedding: NDArray, k: int) -> QueryDocumentsResponse:
|
|
raise NotImplementedError()
|
|
|
|
|
|
@dataclass
|
|
class BankWithIndex:
|
|
bank: MemoryBank
|
|
index: EmbeddingIndex
|
|
|
|
async def insert_documents(
|
|
self,
|
|
documents: List[MemoryBankDocument],
|
|
) -> None:
|
|
model = get_embedding_model()
|
|
for doc in documents:
|
|
content = await content_from_doc(doc)
|
|
chunks = make_overlapped_chunks(
|
|
doc.document_id,
|
|
content,
|
|
self.bank.config.chunk_size_in_tokens,
|
|
self.bank.config.overlap_size_in_tokens
|
|
or (self.bank.config.chunk_size_in_tokens // 4),
|
|
)
|
|
embeddings = model.encode([x.content for x in chunks]).astype(np.float32)
|
|
|
|
await self.index.add_chunks(chunks, embeddings)
|
|
|
|
async def query_documents(
|
|
self,
|
|
query: InterleavedTextMedia,
|
|
params: Optional[Dict[str, Any]] = None,
|
|
) -> QueryDocumentsResponse:
|
|
if params is None:
|
|
params = {}
|
|
k = params.get("max_chunks", 3)
|
|
|
|
def _process(c) -> str:
|
|
if isinstance(c, str):
|
|
return c
|
|
else:
|
|
return "<media>"
|
|
|
|
if isinstance(query, list):
|
|
query_str = " ".join([_process(c) for c in query])
|
|
else:
|
|
query_str = _process(query)
|
|
|
|
model = get_embedding_model()
|
|
query_vector = model.encode([query_str])[0].astype(np.float32)
|
|
return await self.index.query(query_vector, k)
|