Creación de RAG con Milvus, vLLM y Llama 3.1
La Universidad de California - Berkeley donó vLLM, una biblioteca rápida y fácil de usar para la inferencia y el servicio de LLM, a la LF AI & Data Foundation como proyecto en fase de incubación en julio de 2024. Como proyecto miembro, nos gustaría dar la bienvenida a vLLM a la familia de LF AI & Data. 🎉
Los grandes modelos lingüísticos(LLM) y las bases de datos vectoriales suelen emparejarse para construir la Generación Aumentada de Recuperación(RAG), una arquitectura de aplicación de IA popular para abordar las alucinaciones de la IA. Este blog le mostrará cómo construir y ejecutar una RAG con Milvus, vLLM y Llama 3.1. Más concretamente, le mostraré cómo incrustar y almacenar información de texto como incrustaciones vectoriales en Milvus y utilizar este almacén vectorial como base de conocimiento para recuperar de forma eficiente trozos de texto relevantes para las preguntas de los usuarios. Por último, aprovecharemos vLLM para servir el modelo Llama 3.1-8B de Meta para generar respuestas aumentadas por el texto recuperado. ¡Vamos allá!
Introducción a Milvus, vLLM y Meta's Llama 3.1
Base de datos vectorial Milvus
Milvus es una base de datos vectorial distribuida, de código abierto y creada específicamente para almacenar, indexar y buscar vectores para cargas de trabajo de IA Generativa (GenAI). Su capacidad para realizar búsquedas híbridas, filtrado de metadatos, reordenación y gestión eficiente de billones de vectores convierte a Milvus en la elección perfecta para cargas de trabajo de IA y aprendizaje automático. Milvus puede ejecutarse localmente, en un clúster o alojarse en la nube Zilliz totalmente gestionada.
vLLM
vLLM es un proyecto de código abierto iniciado en UC Berkeley SkyLab centrado en optimizar el rendimiento de los servicios LLM. Utiliza una gestión eficiente de la memoria con PagedAttention, batching continuo y kernels CUDA optimizados. En comparación con los métodos tradicionales, vLLM mejora el rendimiento del servicio hasta 24 veces y reduce el uso de memoria de la GPU a la mitad.
Según el artículo "Efficient Memory Management for Large Language Model Serving with PagedAttention", la caché de KV utiliza alrededor del 30% de la memoria de la GPU, lo que puede provocar problemas de memoria. La caché KV se almacena en memoria contigua, pero el cambio de tamaño puede provocar la fragmentación de la memoria, lo que resulta ineficiente para el cálculo.
Imagen 1. Gestión de la memoria caché KV en los sistemas actuales ( documento 2023 Paged Attention)
Al utilizar memoria virtual para la caché KV, vLLM sólo asigna memoria física de la GPU cuando es necesario, lo que elimina la fragmentación de la memoria y evita la preasignación. En las pruebas, vLLM superó a HuggingFace Transformers (HF) y Text Generation Inference (TGI), logrando hasta 24 veces más rendimiento que HF y 3,5 veces más que TGI en las GPU NVIDIA A10G y A100.
Imagen 2. vLLM alcanza un rendimiento 8,5x-15x superior al de HF y 3,3x-3,5x superior al de TGI ( blog 2023 de vLLM).
Llama de Meta 3.1
Meta's Llama 3. 1 se anunció el 23 de julio de 2024. El modelo 405B ofrece un rendimiento puntero en varias pruebas de referencia públicas y tiene una ventana de contexto de 128.000 tokens de entrada con varios usos comerciales permitidos. Junto con el modelo de 405.000 millones de parámetros, Meta lanzó una versión actualizada de Llama3 70B (70.000 millones de parámetros) y 8B (8.000 millones de parámetros). Los pesos de los modelos pueden descargarse del sitio web de Meta.
Una de las principales conclusiones fue que el ajuste de los datos generados puede aumentar el rendimiento, pero los ejemplos de mala calidad pueden degradarlo. El equipo de Llama trabajó intensamente para identificar y eliminar estos malos ejemplos utilizando el propio modelo, modelos auxiliares y otras herramientas.
Construir y realizar la recuperación RAG con Milvus
Prepare su conjunto de datos.
Utilicé la documentación oficial de Milvus como conjunto de datos para esta demostración, que descargué y guardé localmente.
from langchain.document_loaders import DirectoryLoader
# Load HTML files already saved in a local directory
path = "../../RAG/rtdocs_new/"
global_pattern = '*.html'
loader = DirectoryLoader(path=path, glob=global_pattern)
docs = loader.load()
# Print num documents and a preview.
print(f"loaded {len(docs)} documents")
print(docs[0].page_content)
pprint.pprint(docs[0].metadata)
loaded 22 documents
Why Milvus Docs Tutorials Tools Blog Community Stars0 Try Managed Milvus FREE Search Home v2.4.x About ...
{'source': 'https://milvus.io/docs/quickstart.md'}
Descargue un modelo de incrustación.
A continuación, descargue un modelo de incrustación gratuito y de código abierto de HuggingFace.
import torch
from sentence_transformers import SentenceTransformer
# Initialize torch settings for device-agnostic code.
N_GPU = torch.cuda.device_count()
DEVICE = torch.device('cuda:N_GPU' if torch.cuda.is_available() else 'cpu')
# Download the model from huggingface model hub.
model_name = "BAAI/bge-large-en-v1.5"
encoder = SentenceTransformer(model_name, device=DEVICE)
# Get the model parameters and save for later.
EMBEDDING_DIM = encoder.get_sentence_embedding_dimension()
MAX_SEQ_LENGTH_IN_TOKENS = encoder.get_max_seq_length()
# Inspect model parameters.
print(f"model_name: {model_name}")
print(f"EMBEDDING_DIM: {EMBEDDING_DIM}")
print(f"MAX_SEQ_LENGTH: {MAX_SEQ_LENGTH}")
model_name: BAAI/bge-large-en-v1.5
EMBEDDING_DIM: 1024
MAX_SEQ_LENGTH: 512
Trocea y codifica tus datos personalizados como vectores.
Yo utilizaré una longitud fija de 512 caracteres con un solapamiento del 10%.
from langchain.text_splitter import RecursiveCharacterTextSplitter
CHUNK_SIZE = 512
chunk_overlap = np.round(CHUNK_SIZE * 0.10, 0)
print(f"chunk_size: {CHUNK_SIZE}, chunk_overlap: {chunk_overlap}")
# Define the splitter.
child_splitter = RecursiveCharacterTextSplitter(
chunk_size=CHUNK_SIZE,
chunk_overlap=chunk_overlap)
# Chunk the docs.
chunks = child_splitter.split_documents(docs)
print(f"{len(docs)} docs split into {len(chunks)} child documents.")
# Encoder input is doc.page_content as strings.
list_of_strings = [doc.page_content for doc in chunks if hasattr(doc, 'page_content')]
# Embedding inference using HuggingFace encoder.
embeddings = torch.tensor(encoder.encode(list_of_strings))
# Normalize the embeddings.
embeddings = np.array(embeddings / np.linalg.norm(embeddings))
# Milvus expects a list of `numpy.ndarray` of `numpy.float32` numbers.
converted_values = list(map(np.float32, embeddings))
# Create dict_list for Milvus insertion.
dict_list = []
for chunk, vector in zip(chunks, converted_values):
# Assemble embedding vector, original text chunk, metadata.
chunk_dict = {
'chunk': chunk.page_content,
'source': chunk.metadata.get('source', ""),
'vector': vector,
}
dict_list.append(chunk_dict)
chunk_size: 512, chunk_overlap: 51.0
22 docs split into 355 child documents.
Guarde los vectores en Milvus.
Ingest the encoded vector embedding in the Milvus vector database.
# Connect a client to the Milvus Lite server.
from pymilvus import MilvusClient
mc = MilvusClient("milvus_demo.db")
# Create a collection with flexible schema and AUTOINDEX.
COLLECTION_NAME = "MilvusDocs"
mc.create_collection(COLLECTION_NAME,
EMBEDDING_DIM,
consistency_level="Eventually",
auto_id=True,
overwrite=True)
# Insert data into the Milvus collection.
print("Start inserting entities")
start_time = time.time()
mc.insert(
COLLECTION_NAME,
data=dict_list,
progress_bar=True)
end_time = time.time()
print(f"Milvus insert time for {len(dict_list)} vectors: ", end="")
print(f"{round(end_time - start_time, 2)} seconds")
Start inserting entities
Milvus insert time for 355 vectors: 0.2 seconds
Realice una búsqueda de vectores.
Formule una pregunta y busque los trozos vecinos más cercanos de su base de conocimientos en Milvus.
SAMPLE_QUESTION = "What do the parameters for HNSW mean?"
# Embed the question using the same encoder.
query_embeddings = torch.tensor(encoder.encode(SAMPLE_QUESTION))
# Normalize embeddings to unit length.
query_embeddings = F.normalize(query_embeddings, p=2, dim=1)
# Convert the embeddings to list of list of np.float32.
query_embeddings = list(map(np.float32, query_embeddings))
# Define metadata fields you can filter on.
OUTPUT_FIELDS = list(dict_list[0].keys())
OUTPUT_FIELDS.remove('vector')
# Define how many top-k results you want to retrieve.
TOP_K = 2
# Run semantic vector search using your query and the vector database.
results = mc.search(
COLLECTION_NAME,
data=query_embeddings,
output_fields=OUTPUT_FIELDS,
limit=TOP_K,
consistency_level="Eventually")
El resultado recuperado es el que se muestra a continuación.
Retrieved result #1
distance = 0.7001987099647522
('Chunk text: layer, finds the node closest to the target in this layer, and'
...
'outgoing')
source: https://milvus.io/docs/index.md
Retrieved result #2
distance = 0.6953287124633789
('Chunk text: this value can improve recall rate at the cost of increased'
...
'to the target')
source: https://milvus.io/docs/index.md
Construir y realizar la generación RAG con vLLM y Llama 3.1-8B
Instalar vLLM y los modelos de HuggingFace
vLLM descarga por defecto grandes modelos lingüísticos de HuggingFace. En general, siempre que desee utilizar un modelo nuevo en HuggingFace, debe hacer un pip install --upgrade o -U. Además, necesitarás una GPU para ejecutar la inferencia de los modelos Llama 3.1 de Meta con vLLM.
Para obtener una lista completa de todos los modelos compatibles con vLLM, consulte esta página de documentación.
# (Recommended) Create a new conda environment.
conda create -n myenv python=3.11 -y
conda activate myenv
# Install vLLM with CUDA 12.1.
pip install -U vllm transformers torch
import vllm, torch
from vllm import LLM, SamplingParams
# Clear the GPU memory cache.
torch.cuda.empty_cache()
# Check the GPU.
!nvidia-smi
Para saber más sobre cómo instalar vLLM, consulta su página de instalación.
Consigue un token de HuggingFace.
Algunos modelos en HuggingFace, como Meta Llama 3.1, requieren que el usuario acepte su licencia antes de poder descargar los pesos. Por lo tanto, debes crear una cuenta en HuggingFace, aceptar la licencia del modelo y generar un token.
Al visitar esta página de Llama3.1 en HuggingFace, aparecerá un mensaje pidiéndote que aceptes los términos. Haz clic en "Aceptar licencia" para aceptar los términos de Meta antes de descargar los pesos del modelo. La aprobación suele tardar menos de un día.
Tras recibir la aprobación, deberás generar un nuevo token de HuggingFace. Sus antiguos tokens no funcionarán con los nuevos permisos.
Antes de instalar vLLM, inicie sesión en HuggingFace con su nuevo token. A continuación, he utilizado Colab secrets para almacenar el token.
# Login to HuggingFace using your new token.
from huggingface_hub import login
from google.colab import userdata
hf_token = userdata.get('HF_TOKEN')
login(token = hf_token, add_to_git_credential=True)
Ejecutar la generación RAG
En la demostración, ejecutamos el modelo Llama-3.1-8B, que requiere GPU y memoria considerable para girar. El siguiente ejemplo se ejecutó en Google Colab Pro (10$/mes) con una GPU A100. Para obtener más información sobre cómo ejecutar vLLM, puedes consultar la documentación de inicio rápido.
# 1. Choose a model
MODELTORUN = "meta-llama/Meta-Llama-3.1-8B-Instruct"
# 2. Clear the GPU memory cache, you're going to need it all!
torch.cuda.empty_cache()
# 3. Instantiate a vLLM model instance.
llm = LLM(model=MODELTORUN,
enforce_eager=True,
dtype=torch.bfloat16,
gpu_memory_utilization=0.5,
max_model_len=1000,
seed=415,
max_num_batched_tokens=3000)
Escriba una pregunta utilizando contextos y fuentes recuperados de Milvus.
# Separate all the context together by space.
contexts_combined = ' '.join(contexts)
# Lance Martin, LangChain, says put the best contexts at the end.
contexts_combined = ' '.join(reversed(contexts))
# Separate all the unique sources together by comma.
source_combined = ' '.join(reversed(list(dict.fromkeys(sources))))
SYSTEM_PROMPT = f"""First, check if the provided Context is relevant to
the user's question. Second, only if the provided Context is strongly relevant, answer the question using the Context. Otherwise, if the Context is not strongly relevant, answer the question without using the Context.
Be clear, concise, relevant. Answer clearly, in fewer than 2 sentences.
Grounding sources: {source_combined}
Context: {contexts_combined}
User's question: {SAMPLE_QUESTION}
"""
prompts = [SYSTEM_PROMPT]
Ahora, genera una respuesta utilizando los trozos recuperados y la pregunta original introducida en el prompt.
# Sampling parameters
sampling_params = SamplingParams(temperature=0.2, top_p=0.95)
# Invoke the vLLM model.
outputs = llm.generate(prompts, sampling_params)
# Print the outputs.
for output in outputs:
prompt = output.prompt
generated_text = output.outputs[0].text
# !r calls repr(), which prints a string inside quotes.
print()
print(f"Question: {SAMPLE_QUESTION!r}")
pprint.pprint(f"Generated text: {generated_text!r}")
Question: 'What do the parameters for HNSW MEAN!?'
Generated text: 'Answer: The parameters for HNSW (Hiera(rchical Navigable Small World Graph) are: '
'* M: The maximum degree of nodes on each layer oof the graph, which can improve '
'recall rate at the cost of increased search time. * efConstruction and ef: '
'These parameters specify a search range when building or searching an index.'
La respuesta de arriba me parece perfecta.
Si te interesa esta demostración, no dudes en probarla y contarnos lo que piensas. También puedes unirte a nuestra comunidad Milvus en Discord para conversar directamente con todos los desarrolladores de GenAI.
Referencias
Presentación de 2023 vLLM en Ray Summit
Blog de vLLM: vLLM: Servicio LLM fácil, rápido y barato con PagedAttention
Blog útil sobre el funcionamiento del servidor vLLM: Despliegue de vLLM: guía paso a paso
