|
|
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 |
|
|
|
