| import spaces |
| from langchain_text_splitters.character import RecursiveCharacterTextSplitter |
| from sentence_transformers import SentenceTransformer |
|
|
| from src.utils.constants import EMBEDDING_MODEL_NAME |
|
|
| embedding_model = SentenceTransformer(EMBEDDING_MODEL_NAME) |
|
|
|
|
| @spaces.GPU(duration=10) |
| def embed_splitted_docs(splitted_docs): |
| """ |
| Encode the given list of documents using the specified embedding model. |
| |
| Parameters |
| ---------- |
| splitted_docs : List of str |
| The list of documents to be embedded. |
| |
| Returns |
| ------- |
| embeddings : List of numpy.ndarray |
| The embeddings of the given documents. |
| """ |
| embeddings = embedding_model.encode(splitted_docs, show_progress_bar=True) |
| return embeddings |
|
|
|
|
| def split_corpus(corpus, chunk_size=1000): |
| """ |
| Split a given corpus into chunks of a given size. |
| |
| Parameters |
| ---------- |
| corpus : List of str |
| The corpus to be split. |
| chunk_size : int, default=1000 |
| The size of the chunks to be split from the corpus. |
| |
| Returns |
| ------- |
| List of str |
| The list of chunks (splitted documents) from the corpus. |
| """ |
| splitter = RecursiveCharacterTextSplitter( |
| chunk_size=chunk_size, chunk_overlap=0, add_start_index=True |
| ) |
| splitted_docs = splitter.create_documents(corpus) |
| splitted_docs = list(map(lambda x: x.page_content, splitted_docs)) |
|
|
| return splitted_docs |
|
|
