Costruire RAG con Milvus, vLLM e Llama 3.1
L'Università della California - Berkeley ha donato vLLM, una libreria veloce e facile da usare per l'inferenza e il servizio di LLM, alla LF AI & Data Foundation come progetto in fase di incubazione nel luglio 2024. In qualità di progetto membro, diamo il benvenuto a vLLM nella famiglia di LF AI & Data! 🎉
I modelli linguistici di grandi dimensioni(LLM) e i database vettoriali sono solitamente abbinati per costruire la Retrieval Augmented Generation(RAG), una popolare architettura applicativa di AI per affrontare le allucinazioni dell'AI. Questo blog vi mostrerà come costruire ed eseguire una RAG con Milvus, vLLM e Llama 3.1. In particolare, vi mostrerò come incorporare e memorizzare le informazioni di testo come embeddings vettoriali in Milvus e utilizzare questo archivio vettoriale come base di conoscenza per recuperare in modo efficiente i pezzi di testo rilevanti per le domande degli utenti. Infine, sfrutteremo vLLM per utilizzare il modello Llama 3.1-8B di Meta per generare risposte aumentate dal testo recuperato. Immergiamoci!
Introduzione a Milvus, vLLM e Meta Llama 3.1
Il database vettoriale Milvus
Milvus è un database vettoriale distribuito, open-source e appositamente costruito per archiviare, indicizzare e ricercare vettori per carichi di lavoro di IA generativa (GenAI). La sua capacità di eseguire ricerche ibride, filtraggio dei metadati, reranking e di gestire in modo efficiente trilioni di vettori rende Milvus una scelta obbligata per i carichi di lavoro di AI e apprendimento automatico. Milvus può essere eseguito localmente, su un cluster o ospitato nel cloud Zilliz completamente gestito.
vLLM
vLLM è un progetto open-source avviato presso lo SkyLab della UC Berkeley e incentrato sull'ottimizzazione delle prestazioni dei servizi LLM. Utilizza una gestione efficiente della memoria con PagedAttention, batching continuo e kernel CUDA ottimizzati. Rispetto ai metodi tradizionali, vLLM migliora le prestazioni di servizio fino a 24 volte e dimezza l'utilizzo della memoria della GPU.
Secondo il documento "Efficient Memory Management for Large Language Model Serving with PagedAttention", la cache KV utilizza circa il 30% della memoria della GPU, causando potenziali problemi di memoria. La cache KV è archiviata in memoria contigua, ma la modifica delle dimensioni può causare la frammentazione della memoria, che è inefficiente per la computazione.
Immagine 1. Gestione della memoria cache KV nei sistemi esistenti (2023 Paged Attention paper)
Utilizzando la memoria virtuale per la cache KV, vLLM alloca la memoria fisica della GPU solo quando necessario, eliminando la frammentazione della memoria ed evitando la pre-allocazione. Nei test, vLLM ha superato HuggingFace Transformers (HF ) e Text Generation Inference (TGI), ottenendo un throughput fino a 24 volte superiore rispetto a HF e 3,5 volte superiore rispetto a TGI sulle GPU NVIDIA A10G e A100.
Immagine 2. Il throughput del servizio quando ogni richiesta richiede il completamento di tre output paralleli. vLLM raggiunge un throughput 8,5x-15x superiore a HF e 3,3x-3,5x superiore a TGI (2023 vLLM blog).
Meta's Llama 3.1
Meta's Llama 3.1 è stato annunciato il 23 luglio 2024. Il modello 405B offre prestazioni all'avanguardia su diversi benchmark pubblici e ha una finestra di contesto di 128.000 token in ingresso, con vari usi commerciali consentiti. Oltre al modello da 405 miliardi di parametri, Meta ha rilasciato una versione aggiornata di Llama3 70B (70 miliardi di parametri) e 8B (8 miliardi di parametri). I pesi del modello sono disponibili per il download sul sito web di Meta.
Un'intuizione chiave è stata che la messa a punto dei dati generati può aumentare le prestazioni, ma gli esempi di scarsa qualità possono degradarle. Il team di Llama ha lavorato a lungo per identificare e rimuovere questi cattivi esempi utilizzando il modello stesso, modelli ausiliari e altri strumenti.
Costruire ed eseguire il RAG-Retrieval con Milvus
Preparate il vostro set di dati.
Per questa dimostrazione ho utilizzato la documentazione ufficiale di Milvus come set di dati, che ho scaricato e salvato in locale.
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'}
Scaricare un modello di incorporazione.
Successivamente, scaricare un modello di incorporamento gratuito e open source da 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
Tagliare e codificare i dati personalizzati come vettori.
Utilizzerò una lunghezza fissa di 512 caratteri con una sovrapposizione 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.
Salvare i vettori in Milvus.
Inserire l'incorporazione vettoriale codificata nel database dei vettori di 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
Eseguire una ricerca di vettori.
Porre una domanda e cercare i chunk più vicini dalla propria base di conoscenze in 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")
Il risultato recuperato è quello mostrato di seguito.
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
Creare ed eseguire la generazione di RAG con vLLM e Llama 3.1-8B
Installare vLLM e i modelli di HuggingFace
Per impostazione predefinita, vLLM scarica modelli linguistici di grandi dimensioni da HuggingFace. In generale, ogni volta che si vuole usare un modello nuovo su HuggingFace, si deve fare un pip install --upgrade o -U. Inoltre, è necessaria una GPU per eseguire l'inferenza dei modelli Llama 3.1 di Meta con vLLM.
Per un elenco completo di tutti i modelli supportati da vLLM, consultare questa pagina di documentazione.
# (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
Per saperne di più su come installare vLLM, vedere la pagina di installazione.
Ottenere un token HuggingFace.
Alcuni modelli su HuggingFace, come Meta Llama 3.1, richiedono all'utente di accettare la licenza prima di poter scaricare i pesi. Pertanto, è necessario creare un account HuggingFace, accettare la licenza del modello e generare un token.
Quando si visita questa pagina di Llama3.1 su HuggingFace, viene visualizzato un messaggio che chiede di accettare i termini. Fare clic su "Accetta licenza" per accettare i termini di Meta prima di scaricare i pesi del modello. L'approvazione di solito richiede meno di un giorno.
Dopo aver ricevuto l'approvazione, è necessario generare un nuovo token HuggingFace. I vecchi token non funzioneranno con le nuove autorizzazioni.
Prima di installare vLLM, accedere a HuggingFace con il nuovo token. Di seguito, ho utilizzato i segreti di Colab per memorizzare il 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)
Eseguire la generazione RAG
Nella demo, viene eseguito il modello Llama-3.1-8B, che richiede una GPU e una notevole quantità di memoria per essere eseguito. Il seguente esempio è stato eseguito su Google Colab Pro ($10/mese) con una GPU A100. Per saperne di più su come eseguire vLLM, è possibile consultare la documentazione Quickstart.
# 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)
Scrivere un prompt utilizzando i contesti e le fonti recuperate da 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]
Ora, generare una risposta usando i pezzi recuperati e la domanda originale inserita nel 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 risposta qui sopra mi sembra perfetta!
Se siete interessati a questa demo, provatela voi stessi e fateci sapere cosa ne pensate. Siete anche invitati a unirvi alla nostra comunità Milvus su Discord per conversare direttamente con tutti gli sviluppatori di GenAI.
Riferimenti
Presentazione di vLLM 2023 al Ray Summit
Blog vLLM: vLLM: servizio LLM facile, veloce ed economico con PagedAttention
Blog utile sull'esecuzione del server vLLM: Distribuzione di vLLM: una guida passo a passo
