Construindo RAG com Milvus, vLLM e Llama 3.1
A Universidade da Califórnia - Berkeley doou vLLM, uma biblioteca rápida e fácil de usar para inferência e serviço LLM, para a LF AI & Data Foundation como um projeto em estágio de incubação em julho de 2024. Como um projeto de membro companheiro, gostaríamos de dar as boas-vindas ao vLLM que se junta à família LF AI & Data! 🎉
Modelos de linguagem grande(LLMs) e bancos de dados vetoriais são geralmente combinados para construir Retrieval Augmented Generation(RAG), uma arquitetura de aplicativo de IA popular para lidar com alucinações de IA. Este blogue irá mostrar-lhe como construir e executar um RAG com Milvus, vLLM e Llama 3.1. Mais especificamente, mostrarei como incorporar e armazenar informações de texto como embeddings vetoriais no Milvus e usar esse armazenamento vetorial como uma base de conhecimento para recuperar com eficiência pedaços de texto relevantes para as perguntas do usuário. Por fim, utilizaremos o vLLM para servir o modelo Llama 3.1-8B do Meta para gerar respostas aumentadas pelo texto recuperado. Vamos mergulhar de cabeça!
Introdução ao Milvus, vLLM e Llama 3.1 do Meta
Base de dados vetorial Milvus
O Milvus é uma base de dados de vectores distribuída, de código aberto, criada propositadamente para armazenar, indexar e pesquisar vectores para cargas de trabalho de IA generativa (GenAI). Sua capacidade de realizar pesquisa híbrida, filtragem de metadados, reranking e lidar com trilhões de vetores de forma eficiente faz do Milvus uma escolha para cargas de trabalho de IA e aprendizado de máquina. O Milvus pode ser executado localmente, num cluster ou alojado no Zilliz Cloud totalmente gerido.
vLLM
O vLLM é um projeto de código aberto iniciado no UC Berkeley SkyLab focado na otimização do desempenho do serviço LLM. Ele usa gerenciamento eficiente de memória com PagedAttention, batching contínuo e kernels CUDA otimizados. Em comparação com os métodos tradicionais, o vLLM melhora o desempenho de serviço em até 24x, reduzindo o uso de memória da GPU pela metade.
De acordo com o documento "Efficient Memory Management for Large Language Model Serving with PagedAttention", o cache KV usa cerca de 30% da memória da GPU, levando a possíveis problemas de memória. A cache KV é armazenada em memória contígua, mas a alteração do tamanho pode causar fragmentação da memória, o que é ineficiente para a computação.
Imagem 1. Gestão da memória cache KV nos sistemas existentes (2023 Paged Attention paper)
Ao usar a memória virtual para o cache KV, o vLLM aloca apenas a memória física da GPU conforme necessário, eliminando a fragmentação da memória e evitando a pré-alocação. Nos testes, o vLLM superou o HuggingFace Transformers (HF) e o Text Generation Inference (TGI), alcançando uma taxa de transferência até 24 vezes maior que o HF e 3,5x maior que o TGI nas GPUs NVIDIA A10G e A100.
Imagem 2. O vLLM atinge uma taxa de transferência 8,5x-15x superior à do HF e 3,3x-3,5x superior à do TGI ( blogue 2023 vLLM).
Meta's Llama 3.1
O Llama 3.1 da Meta foi anunciado em 23 de julho de 2024. O modelo 405B oferece desempenho de ponta em vários benchmarks públicos e tem uma janela de contexto de 128.000 tokens de entrada com várias utilizações comerciais permitidas. Juntamente com o modelo de 405 mil milhões de parâmetros, a Meta lançou uma versão actualizada do Llama3 70B (70 mil milhões de parâmetros) e 8B (8 mil milhões de parâmetros). Os pesos dos modelos estão disponíveis para descarregamento no sítio Web da Meta.
Uma das principais conclusões foi que o ajuste fino dos dados gerados pode aumentar o desempenho, mas exemplos de baixa qualidade podem degradá-lo. A equipa do Llama trabalhou extensivamente para identificar e remover estes maus exemplos utilizando o próprio modelo, modelos auxiliares e outras ferramentas.
Construir e executar o RAG-Retrieval com o Milvus
Prepare o seu conjunto de dados.
Utilizei a documentação oficial do Milvus como conjunto de dados para esta demonstração, que descarreguei e guardei 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'}
Descarregue um modelo de incorporação.
Em seguida, descarregue um modelo de incorporação gratuito e de código aberto do 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
Divida e codifique seus dados personalizados como vetores.
Vou usar um comprimento fixo de 512 caracteres com 10% de sobreposição.
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 os vectores em Milvus.
Ingerir a incorporação do vetor codificado na base de dados de vectores do Milvus.
# 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
Efetuar uma pesquisa de vectores.
Faça uma pergunta e procure os pedaços vizinhos mais próximos da sua base de conhecimentos em 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")
O resultado obtido é o seguinte.
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 e efetuar a geração de RAG com vLLM e Llama 3.1-8B
Instalar o vLLM e os modelos do HuggingFace
O vLLM baixa modelos de linguagem grandes do HuggingFace por padrão. Em geral, sempre que você quiser usar um novo modelo no HuggingFace, você deve fazer um pip install --upgrade ou -U. Além disso, você precisará de uma GPU para executar a inferência dos modelos Llama 3.1 do Meta com vLLM.
Para obter uma lista completa de todos os modelos suportados pelo vLLM, consulte esta página de documentação.
# (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 mais sobre como instalar o vLLM, consulte sua página de instalação.
Obter um token HuggingFace.
Alguns modelos no HuggingFace, como o Meta Llama 3.1, exigem que o usuário aceite sua licença antes de poder baixar os pesos. Portanto, é necessário criar uma conta no HuggingFace, aceitar a licença do modelo e gerar um token.
Ao visitar esta página do Llama3.1 no HuggingFace, aparecerá uma mensagem pedindo para concordar com os termos. Clique em "Accept License" para aceitar os termos do Meta antes de descarregar os pesos do modelo. A aprovação geralmente leva menos de um dia.
Depois de receber a aprovação, deve gerar um novo token HuggingFace. Seus tokens antigos não funcionarão com as novas permissões.
Antes de instalar o vLLM, faça login no HuggingFace com seu novo token. Abaixo, eu usei Colab secrets para armazenar o 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)
Executar a geração de RAG
Na demonstração, nós rodamos o modelo Llama-3.1-8B, que requer GPU e memória considerável para rodar. O exemplo a seguir foi executado no Google Colab Pro ($10/mês) com uma GPU A100. Para saber mais sobre como executar o vLLM, consulte a documentação de início 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)
Escreva um prompt usando contextos e fontes recuperados do 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]
Agora, gere uma resposta usando os pedaços recuperados e a pergunta original inserida no 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.'
A resposta acima parece-me perfeita!
Se estiver interessado nesta demonstração, não hesite em experimentá-la e diga-nos o que pensa. Também pode juntar-se à nossa comunidade Milvus no Discord para conversar diretamente com todos os programadores do GenAI.
Referências
Documentação oficial do vLLM e página do modelo.
Apresentação do vLLM 2023 na Ray Summit
Blog do vLLM: vLLM: Serviço LLM fácil, rápido e barato com PagedAttention
Blog útil sobre a execução do servidor vLLM: Implantando o vLLM: um Guia Passo-a-Passo
