Variables de modelo de Embedding (1024 dimensiones)

emb_model = "dunzhang/stella_en_400M_v5" # Modelo de embedding
emodel = SentenceTransformer(emb_model, trust_remote_code=True, device="cpu", config_kwargs={"use_memory_efficient_attention": False, "unpad_inputs": False})
k_retrieve = 5 # Documentos a retornar con el Semantico

This model supports two prompts: "s2p_query" and "s2s_query" for sentence-to-passage and sentence-to-sentence tasks, respectively.

They are defined in config_sentence_transformers.json

query_prompt_name = "s2p_query"

Embedding

def doc_embeddings(all_md_header_splits):
only_content = [doc.page_content for doc in all_md_header_splits] # Crear una lista con solo el page_content de cada documento
doc_embeddings = emodel.encode(only_content) # Hago los embeddings
return doc_embeddings, only_content # Retorno los vectores de embedding y los chunks

Retrieve

def similarity_search(query, doc_embeddings, only_content): # Busco el contexto con la busqueda de similaridad
query_embeddings = emodel.encode(query, prompt_name=query_prompt_name) # Hago el embedding de la query
all_docs = emodel.similarity(query_embeddings, doc_embeddings) # Hago la busqueda

# Asegurarse de que k_retrieve no sea mayor que la cantidad de documentos disponibles
num_docs = all_docs.numel()                                                      # numel() devuelve el número total de elementos en el tensor
k_safe = min(k_retrieve, num_docs)                                               # Tomar el menor entre k_retrieve y el número de elementos disponibles
top_values, top_indices = torch.topk(all_docs.flatten(), k_safe)                 # Utilizar torch.topk para obtener los K valores solicitados
texto_combinado = ""                                                             # Inicializar una variable para almacenar el texto
retrieved_docs = []                                                              # Lista para devolver con los documentos

for indice in top_indices:                                                       # Iterar sobre los índices en top_indices
    indice_como_numero = indice.item()                                           # Tomo el valor del tensor para usarlo como indice
    texto_combinado += only_content[indice_como_numero] + " "                    # Concatenamos el texto de esa posición
    retrieved_docs.append(only_content[indice_como_numero])                      # Añadimos el contenido del documento a la lista
    
context = "Contexto: " + texto_combinado.strip()                                 # Eliminamos los espacios al final
return context, retrieved_docs                                                   # Retorno el contexto en texto puro y en formado de listado de documentos