Retrieval-Augmented Generation (RAG) LLMs provide organization the ability to deploy LLMs with guide rails to prevent LLM hallucination. The following guides demonstrate:

• RAG LLM Prerequisites
• RAG LLM Deployment and example inferences

The following LLM models have been tested with Wallaroo, with the majority packaged as BYOP models.

• llama v2 7B standard on 1 GPU
• llama v2 7B chat on 1 GPU
• llama v2 7B instruct on 1 GPU
• llama v2 70B quantized on 1 GPU
• llama v3 8B standard on 1 GPU
• llama v2 8B instruct on 1 GPU
• llama v2 7B Chat quantized with llamacpp on ARM and X86
• llama v3 8B instruct quantized with llamacpp on ARM and X86

For access to these sample models and a demonstration on using LLMs with Wallaroo:

• Contact your Wallaroo Support Representative OR
• Schedule Your Wallaroo.AI Demo Today

Prerequisites

Before inferencing with the RAG LLM, the following prerequisites must be met.

• Vector indexed database: RAG LLM’s require that the vector indexed database, such as MongoDB Atlas Vector Search.
• Embedder model: This processes the incoming text into a vector used by the RAG LLM to query the database for the context. From the context, the RAG LLM is able to build the response without hallucinations.

Retrieval-Augmented Generation (RAG) LLMs deployed to Wallaroo utilize a two-step pipeline, with two models:

• A Wallaroo BYOP framework model that implements the feature extractor to create the embedding from the input text. For this example, the bge-base-en model is used.
• A RAG LLM packaged in a Wallaroo BYOP framework model.

Overview

The RAG LLM process takes the following steps:

• Inference requests first pass through the feature extractor model that outputs the embedding. This is a list of floats that the RAG LLM uses to query the database for its context.
• Both the embedding and the origin input is passed to the RAG LLM.
• The RAG LLM queries the vector database for the context from which to build it’s response. This context prevents hallucinations by providing guide rails that the RAG LLM uses to construct its response.
• Once finished, the response is submitted as the inference output.

Details on the Feature Extractor model and the an example of querying the database by its vector index are available below.

def _predict(self, input_data: InferenceData):
        inputs = input_data["text"].tolist()
        texts = np.array([str(x) for x in input_data["text"]])

        encoded_inputs = self.model["tokenizer"](
            inputs, padding=True, truncation=True, return_tensors="pt"
        )

        with torch.no_grad():
            model_output = self.model["model"](**encoded_inputs)
            sentence_embeddings = model_output[0][:, 0]

        sentence_embeddings = torch.nn.functional.normalize(
            sentence_embeddings, p=2, dim=1
        )

        embeddings = np.array(
            [sentence_embeddings[i].cpu().numpy() for i in range(len(inputs))]
        )

        return {"embedding": embeddings, "text": texts}

def _predict(self, input_data: InferenceData):
        db = client.sample_mflix
        collection = db.movies

        generated_texts = []
        prompts = input_data["text"].tolist()
        embeddings = input_data["embedding"].tolist()

        for prompt, embedding in zip(prompts, embeddings):
            query_results = collection.aggregate(
                [
                    {
                        "$vectorSearch": {
                            "queryVector": embedding,
                            "path": "plot_embedding_hf",
                            "numCandidates": 50,
                            "limit": 10,
                            "index": "PlotSemanticSearch",
                        }
                    }
                ]
            )

            context = " ".join([result["plot"] for result in query_results])

            result = self.model(
                f"Q: {prompt} C: {context} A: ",
                max_tokens=512,
                stop=["Q:", "\n"],
                echo=False,
            )
            generated_texts.append(result["choices"][0]["text"])

        return {"generated_text": np.array(generated_texts)}

Upload the Models

Both the feature extractor model and the RAG LLM are uploaded via either the Wallaroo SDK or the Wallaroo MLOps API.

Both models are packaged in the Wallaroo BYOP framework.

import pyarrow as pa

input_schema = pa.schema([
    pa.field('text', pa.string())
])

output_schema = pa.schema([
    pa.field('embedding', 
        pa.list_(
            pa.float64(), list_size=768
        ),
    ),
    pa.field('text', pa.string())
])

import wallaroo

# establish the Wallaroo client

wl = wallaroo.Client()


bge = wl.upload_model('byop-bge-base-pipe-v1', 
    './models/byop_bge_base_pipe.zip',
    framework=Framework.CUSTOM,
    input_schema=input_schema,
    output_schema=output_schema,
)

import pyarrow as pa

input_schema = pa.schema([
    pa.field('text', pa.string()),
    pa.field('embedding', pa.list_(pa.float32(), list_size=768))
]) 

output_schema = pa.schema([
    pa.field('generated_text', pa.string()),
])

rag-llm = wl.upload_model('byop-llamacpp-rag-v1', 
    './models/byop_llamacpp_rag.zip',
    framework=Framework.CUSTOM,
    input_schema=input_schema,
    output_schema=output_schema,
)

export WALLAROO_USERNAME="username"
export WALLAROO_PASSWORD="password"
export WALLAROO_AUTH_URL="keycloak.wallaroo.example.com"

TOKEN=$(curl -s -X POST "https://$WALLAROO_AUTH_URL/auth/realms/master/protocol/openid-connect/token" \
                 -d client_id=sdk-client \
                 -d username=${WALLAROO_USERNAME} \
                 -d password=${WALLAROO_PASSWORD} \
                 -d grant_type=password | jq .access_token )
echo $TOKEN

"abc123"

import pyarrow as pa
import base64

input_schema = pa.schema([
    pa.field('input_1', pa.list_(pa.float32(), list_size=10)),
    pa.field('input_2', pa.list_(pa.float32(), list_size=5))
])
output_schema = pa.schema([
    pa.field('output_1', pa.list_(pa.float32(), list_size=3)),
    pa.field('output_2', pa.list_(pa.float32(), list_size=2))
])

encoded_input_schema = base64.b64encode(
                bytes(input_schema.serialize())
            ).decode("utf8")

encoded_output_schema = base64.b64encode(
                bytes(output_schema.serialize())
            ).decode("utf8")

curl --progress-bar -X POST \
  -H "Content-Type: multipart/form-data" \
  -H "Authorization: Bearer $TOKEN" \
  -F 'metadata=\ 
      { \
        "name": $NAME, \
        "visibility": "private", \
        "workspace_id": $WORKSPACE_ID, \
        "conversion": \
          {\
            "framework": $FRAMEWORK, \
            "python_version": "3.8", \
            "requirements": []}, \
            "input_schema": $INPUT_SCHEMA, \
            "output_schema": $OUTPUT_SCHEMA, \
            "arch": 'x86' \
          };type=application/json' \
  -F "file=@$MODEL_PATH;type=application/octet-stream" \
$URL/v1/api/models/upload_and_convert | cat

RAG LLM Inference Example

The following example assumes that the two models are already uploaded and saved to the following variables:

• bge: The Feature Extractor that generates the embedding for the RAG LLM.
• rag-llm: The RAG LLM that uses the embedding to query the vector database index, and uses that result as the context to generate the text.

With the models uploaded, they are deployed in a Wallaroo pipeline through the following process:

• Define the deployment configuration: This sets what resources are applied to each model on deployment. For more details, see Deployment Configuration.
• Add the feature extractor model and RAG LLM as model steps: This sets the structure where the feature extractor model converts the request to a vector, which is used as the input by the RAG LLM to generate the final response.
• Deploy the models: This step allocates resources to the feature extractor and LLM. At this point, the models are ready for inference requests.

Set the Deployment Configuration

The deployment configuration sets the following resources to the models:

• Feature extractor
  • 4 CPUs
  • 3 Gi RAM
• RAG LLM
  • 4 CPUs
  • 6 Gi RAM

This is represented in the following code.

deployment_config = DeploymentConfigBuilder() \
    .cpus(1).memory('2Gi') \
    .sidekick_cpus(bge, 4) \
    .sidekick_memory(bge, '3Gi') \
    .sidekick_cpus(rag-llm, 4) \
    .sidekick_memory(rag-llm, '6Gi') \
    .build()

Deployment configurations are adjusted as required based on attributes including model size, throughput, latency, and performance requirements. Deployment configuration changes do not impact Wallaroo Inference endpoints (including name, url, etc), providing no interruption for production deployments.

Set the Model Steps

In this step, we add the feature extractor model and the RAG LLM as pipeline steps.

We create the pipeline with the wallaroo.client.Client.build_pipeline, then add each model as pipeline steps with the feature extractor as the first step with the wallaroo.pipeline.Pipeline.add_model_step method. This sets the stage for the feature extractor model to provide its outputs as the inputs for the RAG LLM.

pipeline = wl.build_pipeline("byop-rag-llm-bge-v1")
pipeline.add_model_step(bge)
pipeline.add_model_step(rag-llm)

Deploy the RAG LLM

Everything is now set and we deploy the models through the wallaroo.pipeline.Pipeline.deploy(deployment_config) method, providing the deployment configuration we set earlier. This assigns the resources from the cluster to the model’s exclusive use.

Once the deployment is complete, the RAG LLM is ready for inference requests.

pipeline.deploy(deployment_config=deployment_config)

Inference Example

Inference requests are submitted either as pandas DataFrames or Apache Arrow tables. The following example shows submitting a pandas DataFrame with the query to suggest an action movie. The response is returned as a pandas DataFrame, and we extract the generated text from there.

data = pd.DataFrame({"text": ["Suggest me an action movie, including it's name"]})

result = pipeline.infer(data)
print(result['out.generated_text'].values[0])

For access to these sample models and a demonstration on using LLMs with Wallaroo:

• Contact your Wallaroo Support Representative OR
• Schedule Your Wallaroo.AI Demo Today

Tutorials

For tutorials on using RAG LLMs and using Wallaroo Inference Automation with the embedding model to populate the vector database, see RAG LLM Tutorials.

FAQs

Why use MongoDB Atlas for RAG-LLMs?

MongoDB’s recently released Atlas Vector Search allows for the integration of vector search capabilities into MongoDB Atlas. This allows for the creation of a vector index on a collection and the ability to perform vector searches on that collection. This is particularly useful for RAG-LLMs as it allows for the integration of vector search capabilities into the llama.cpp pipeline. Wallaroo supports other vector database connections. Please contact us to learn more!