AI Frameworks

LlamaIndex Integration

NeuReality integrates with LlamaIndex easily, both on LLMs and Embedding models.

We have two options for integration:

  1. Integrating our client library with LlamaIndex’s existing classes, allowing requests to be directed to NR1 instead of traditional servers like OpenAI. See Figure 1: Alt text

  2. Utilizing a custom class we’ve developed, currently provided as a wheel package. In the future, we aim to contribute this implementation to LlamaIndex, eliminating the need for a separate wheel. See Figure 2: Alt text

LLM - the first option

Using the OpenAILike class from LlamaIndex, and injecting the NROpenAI as the openai_client (which is one of the OpenAILike class parameters).

NROpenAI expects to get base_url and api_key:

  • base_url: The base url should include the ip, port, namespace and cg dir name of the NR1 and artifacts for the LLM we want to use. This is the needed format of the base_url:

    f"{ip}:{port}/{namespace}/{cg_dir_name}"

  • api_key: Not really needed, so you can always use “NR_API_KEY”

You can get this NROpenAI class by installing the wheel (taken from our artifactory)

export NR_OPEN_AI_WHEEL= # path to NR OpenAI wheel from artifactory
pip install $NR_OPEN_AI_WHEEL --force-reinstall

After installing, you can use it. For example:

Assuming the NR1 we are using is 172.23.8.54, the port is 5200, the namespace is “default” and the dir name of the artifacts is “llama_3_1_8b_instruct_1_128_8064_rag_demo”, meaning:

ip="172.23.8.54"
port="5200"
namespace="default"
cg_dir_name="llama_3_1_8b_instruct_1_128_8064_rag_demo"

You can do this (change the other parameters according to your needs):

from llama_index.llms.openai_like import OpenAILike
from neureality.client.nr_openai_client import NROpenAI # can be done after installing the wheel

# for LLM
llm_base_url = f"{ip}:{port}/{namespace}/{cg_dir_name}" # defined above

client = NROpenAI(api_key = "NR_API_KEY", base_url=llm_base_url)

nr_llm = OpenAILike(
        llm_model_name,
        openai_client=client,
        max_tokens=128,
        temperature=0,
        context_window=8192,
        completion_to_prompt=completion_to_prompt,
)
Settings.llm = nr_llm

Embedding model - the second option

We have created a new class (NREmbedding) based on the BaseEmbedding LlamaIndex class.

NREmbedding Class

The NREmbedding class should be used as an embedding model based on LlamaIndex. The are two main differences between this and the existing embedding classes in LlamaIndex:

  1. The ability to integrate a custom client library into the class, enabling communication tailored to the user’s requirements. In our case, we aim to integrate the NROpenAI client library, which follows the OpenAI API structure but redirects requests to NeuReality servers instead of OpenAI’s servers.

  2. You can use any model you choose. When utilizing the NREmbedding class on NeuReality, the model selection depends on whether it can be compiled for our accelerator, rather than on external support, such as from OpenAI.

How to use

To use the NREmbedding class, you should install a wheel which includes this class, you can download the wheel from our artifactory and install it as follows:

export NR_EMBEDDING_WHEEL= # path to NR Embedding wheel from artifactory
pip install $NR_EMBEDDING_WHEEL --force-reinstall

After installing, you can use it. For example:

from llamaindex_embedding_class.nr_embedding import NREmbedding # can be done after installing the wheel
embedding_model_name="BAAI/bge-large-en-v1.5" # change according to your needs

# These params can be set in env vars instead inserting them as params to the NREmbedding constructor
ip="172.23.8.143"
port="5200"
namespace="default"
cg_dir_name="embedding_rag"

embedding_model = NREmbedding(model_name=embedding_model_name, ip=ip, port=port, namespace=namespace, cg_dir_name=cg_dir_name)

# to use this NREmbedding class using LlamaIndex
from llama_index.core import Settings
Settings.embed_model = embedding_model

The NREmbedding class parameters are:

  • model_name (str) - the model name (since it should be passed to the embedding functions)

  • instruction (str) - instruction sentence (not used for now)

  • ip (str) - the IP of the NR1 board

  • port (str) - the port of the ai-server on the NR1 board. It is defined through the provisioning phase.

  • namespace (str) - the namespace where the artifacts have been provisioned.

  • cg_dir_name (str) - the directory name of the artifacts, which we want to infer (after we have provisoned them).

Notice:

The ip, port, namespace and cg_dir_name can be defined as environment variables as well (instead of passing it to the NREmedding constructor).

The main important functions (the same as the functions of a BaseEmbedding model in LlamaIndex):

  • get_query_embedding(str)

  • get_text_embedding(str)

  • get_text_embeddings(List[str])

LangChain Support

LangChain is a framework designed to simplify the development of applications powered by language models. It provides tools for chaining together multiple components, such as LLMs, embeddings, and vector stores, enabling advanced workflows like Retrieval-Augmented Generation (RAG), question answering, and more.

This guide explains how to use LangChain to run inference for embedding models, LLM models, and Retrieval-Augmented Generation (RAG) on NeuReality’s NR1 Module boards.


Prerequisites

  1. Install the required dependencies:

    LangChain can be integrated with NeuReality’s NR1 boards using NeuReality’s OpenAI-compatible endpoint, which is supported by the LangChain framework as a drop-in replacement for OpenAI’s own service. To ensure compatibility with code presented in this guide, install all required dependencies using the following command:

    pip install pytest numpy langchain openai langchain-openai langchain-community langchain-core langchain-text-splitters langgraph bs4 typing-extensions
    
  2. Ensure you have access to the necessary API keys and endpoints for the OpenAI-compatible endpoints you will use (LangChain itself does not require API keys).

  3. Set up the required environment variables or pass the API keys directly in the code.


Running Embedding Model Inference

Embedding models are designed to convert textual data into numerical vector representations. These vectors capture semantic meaning, enabling efficient comparison, clustering, and retrieval of text data. Embedding models are widely used in applications such as search engines, recommendation systems, and natural language processing tasks.

Embedding models typically accept a single string of text, such as a sentence, paragraph, or document.

The outputs of embedding models are numerical vectors, which can be used at later stages for various tasks. Text can be represented with a high-dimensional vector such as [0.33, -0.14, 0.26, ...].

Key Characteristics of Embedding Outputs:

  1. Dimensionality - The dimensionality of the output vectors depends on the embedding model used. Common dimensions range from hundreds to thousands.

  2. Semantic Similarity - Vectors produced by embedding models are designed to capture semantic similarity. Texts with similar meanings will have vectors that are closer in the vector space.

  3. Normalization - Some embedding models normalize the output vectors to have a unit length, ensuring consistent scaling for similarity calculations.

Example use cases:

  1. Similarity Search - Compare the query embedding with text embeddings to find the most relevant documents.

  2. Clustering - Group similar texts based on their embeddings using clustering algorithms.

  3. Classification - Use embeddings as features for machine learning models to classify text data.

  4. Retrieval-Augmented Generation (RAG) - Combine embeddings with LLMs to retrieve relevant context and generate informed responses.

Embedding models are a foundational component of modern AI frameworks, enabling efficient and meaningful processing of textual data.

Follow these steps to implement embedding model inference with LangChain and the NR1 Module:

  1. Set up the OpenAI client and LangChain embeddings: The following code demonstrates how to retrieve model details using the OpenAI API and create an instance of the OpenAIEmbeddings class. This example assumes the model responds on port 8000 (default). If you encounter connection issues, ensure the model is provisioned with the OpenAI REST API and that the port matches the provisioned model’s configuration.

    from openai import OpenAI
    from langchain_openai import OpenAIEmbeddings
    
    openai_api_base = "http://<ip_nr1_board>:8000/v1"
    client = OpenAI(api_key="<nr_openai_api_key>", base_url=openai_api_base)
    model = client.models.list().data[0].id
    
    embeddings = OpenAIEmbeddings(
        model=model,
        openai_api_key="<nr_openai_api_key>",
        openai_api_base=openai_api_base,
        chunk_size=256,
    )
    
  2. Generate embeddings for a query: Use the embedding instance with the LangChain framework as demonstrated below:

    input_text = "What is the meaning of life?"
    vector = embeddings.embed_query(input_text)
    

    You can also use embeddings for further processing within the LangChain framework. An example of such use is located in the RAG section.


Running LLM Model Inference

Large Language Models (LLMs) are advanced machine learning models trained on vast amounts of textual data to understand, generate, and manipulate human language. These models leverage deep learning techniques, particularly transformer architectures, to process and generate text with remarkable fluency and coherence.

LLMs are designed to perform a wide range of tasks, including text completion, summarization, translation, question answering, and more. Their ability to understand context and generate human-like responses makes them invaluable in various applications.

Key characteristics of LLMs:

  1. Context Awareness: LLMs can understand and generate text based on the context provided, making them suitable for tasks requiring nuanced understanding.

  2. Scalability: They can handle large-scale data and generate responses for complex queries.

  3. Adaptability: LLMs can be fine-tuned for specific domains or tasks, enhancing their performance in specialized applications.

  4. Multilingual Support: Many LLMs are trained on multilingual datasets, enabling them to process and generate text in multiple languages.

Common Use Cases for LLMs:

  1. Text Completion

  2. Summarization

  3. Translation

  4. Question Answering

  5. Creative Writing

  6. Code Generation

  7. Sentiment Analysis

  8. Retrieval-Augmented Generation (RAG)

  9. Personalization

  10. Education and Training

After the model is correctly provisioned with the OpenAI REST API, follow these steps to implement LLM with LangChain framework and the NR1 Module:

  1. Set up the OpenAI client and LangChain LLM: Create an LLM class instance similar to the embedding class using LangchainOpenAI. Ensure the correct IP and port are provided, and verify that the model is provisioned correctly (perform a health check using REST).

    from openai import OpenAI
    from langchain_openai import OpenAI as LangchainOpenAI
    
    openai_api_base = "http://<ip_nr1_board>:8000/v1"
    client = OpenAI(api_key="<nr_openai_api_key>", base_url=openai_api_base)
    model = client.models.list().data[0].id
    
    llm = LangchainOpenAI(
        model=model,
        openai_api_key="<nr_openai_api_key>",
        openai_api_base=openai_api_base,
    )
    
  2. Invoke the LLM for a query: Run inference using the invoke method, which is a blocking operation that waits for the entire response.

    input_text = "The meaning of life is "
    response = llm.invoke(input_text)
    print(input_text + response)
    
  3. Stream responses: Alternatively, generate the LLM response using a non-blocking operation with the stream method.

    for chunk in llm.stream(input_text):
        print(chunk, end="")
    

Running RAG (Retrieval-Augmented Generation)

Retrieval-Augmented Generation (RAG) is a technique that combines the capabilities of large language models (LLMs) with external knowledge sources to improve the accuracy and relevance of generated responses. It is particularly useful for tasks requiring up-to-date or domain-specific information that may not be fully captured within the pre-trained parameters of the language model.

The RAG process typically involves two main steps:

  1. Retrieval: Relevant documents or pieces of information are fetched from an external knowledge base, such as a database, search engine, or vector store. This retrieval step ensures that the model has access to the most pertinent and contextually relevant data for the task at hand.

  2. Generation: The retrieved information is then passed to the language model as context or input, enabling it to generate responses that are informed by the external data. This step leverages the model’s natural language generation capabilities to produce coherent and contextually accurate outputs.

Advantages of RAG:

  • Improved Accuracy: By incorporating external knowledge, RAG reduces the likelihood of hallucinations (i.e., generating incorrect or fabricated information).

  • Domain-Specific Knowledge: It allows the model to access specialized or proprietary information that may not be included in its training data.

  • Dynamic Updates: Since the retrieval step can query live or frequently updated sources, RAG enables responses to reflect the latest information.

Use Cases:

  • Question answering systems that require precise and factual responses.

  • Customer support bots that need access to company-specific documentation.

  • Research assistants that synthesize information from multiple sources.

  • Any application where combining generative AI with external data enhances performance and reliability.

Implementation typically involves integrating a retrieval mechanism (e.g., vector search, semantic search) with a generative model, ensuring seamless interaction between the two components.

⚠️ Warning: NeuReality SW 1.3 requires each model to be provisioned on different ports. If you attempt to provision multiple models on the same port, they will not function correctly. Software versions higher than 1.3 will support multiple models provisioned on the same port.

  1. Set up embedding and LLM models: After the models are correctly provisioned (e.g., on ports 8010 and 8020), create LangChain embeddings and LLM class instances as shown below:

    from openai import OpenAI
    from langchain_openai import OpenAIEmbeddings, OpenAI as LangchainOpenAI
    
    embedding_openai_api_base = "http://<ip_nr1_board>:8010/v1"
    llm_openai_api_base = "http://<ip_nr1_board>:8020/v1"
    
    embeddings = OpenAIEmbeddings(
        model="<embedding_model>",
        openai_api_key="<nr_openai_api_key>",
        openai_api_base=embedding_openai_api_base,
    )
    
    llm = LangchainOpenAI(
        model="<llm_model>",
        openai_api_key="<nr_openai_api_key>",
        openai_api_base=llm_openai_api_base,
    )
    
  2. Load and chunk documents: One of the possible inputs for RAG is a webpage. This example demonstrates how to use a webpage to provide information that enhances the LLM’s context.

    • Replace url with the desired webpage address.

    • Choose chunk_size and chunk_overlap carefully, based on your embedding model’s maximum input size.

    In the example provided, an embedding model capable of handling 256 input tokens was used, and 512 characters were chosen as a safe length to fit within the limit.

    The following code uses webpage data, which is then split into chunks manageable by the embedding model.

    from langchain_community.document_loaders import WebBaseLoader
    from langchain_text_splitters import RecursiveCharacterTextSplitter
    
    loader = WebBaseLoader(web_paths=["<url>"])
    docs = loader.load()
    
    text_splitter = RecursiveCharacterTextSplitter(chunk_size=512, chunk_overlap=128)
    all_splits = text_splitter.split_documents(docs)
    
  3. Index documents in a vector store: Data splits taken from previous steps need to be stored in storage for retrieval, and they have to be assigned embedding values that will be used to provide similarity information to the user’s query.

    LangChain provides an easy-to-use InMemoryVectorStore which will be sufficient for this example (for solutions using a large amount of data, you may need to consider other options supported by LangChain).

    The in-memory vector store requires an embedding class instance to be provided at initialization, which will be used to automatically generate embedding vectors for the provided document splits.

    from langchain_core.vectorstores import InMemoryVectorStore
    
    vector_store = InMemoryVectorStore(embeddings)
    vector_store.add_documents(documents=all_splits)
    
  4. Define the RAG pipeline: To retrieve relevant data, similarity_search is used to extract and construct a prompt for the LLM using LangChain Hub. The constructed prompt is then passed to the LLM model.

    The graph for this process is constructed as shown below.

    from langchain import hub
    from langgraph.graph import START, StateGraph
    from typing_extensions import TypedDict, List
    from langchain_core.documents import Document
    
    prompt = hub.pull("rlm/rag-prompt")
    
    class State(TypedDict):
        question: str
        context: List[Document]
        answer: str
    
    def retrieve(state: State):
        retrieved_docs = vector_store.similarity_search(state["question"])
        return {"context": retrieved_docs}
    
    def generate(state: State):
        docs_content = "\n\n".join(doc.page_content for doc in state["context"])
        messages = prompt.invoke(
            {"question": state["question"], "context": docs_content}
        )
        response = llm.invoke(messages)
        return {"answer": response}
    
    graph_builder = StateGraph(State).add_sequence([retrieve, generate])
    graph_builder.add_edge(START, "retrieve")
    graph = graph_builder.compile()
    
  5. Execute the RAG pipeline: After performing the previous steps, the compiled graph can be executed using the code below.

    response = graph.invoke({"question": "What is Task Decomposition?"})
    print(response["answer"])
    

📜 Full Code Example

Based on the LangChain RAG tutorial on langchain.com.

    ## Setup embedding model
    embedding_openai_api_base = f"http://<ip_nr1_board>:8010/v1"

    client = OpenAI(
        api_key=nr_openai_api_key,
        base_url=embedding_openai_api_base,
    )

    models = client.models.list()
    model = models.data[0].id

    embeddings = OpenAIEmbeddings(
        model=model,
        tiktoken_enabled=False,
        check_embedding_ctx_length=False,
        openai_api_key=nr_openai_api_key,
        openai_api_base=embedding_openai_api_base,
        chunk_size=256,
    )

    ## Setup LLM Model
    llm_openai_api_base = f"http://<ip_nr1_board>:8020/v1"
    client = OpenAI(
        api_key=nr_openai_api_key,
        base_url=llm_openai_api_base,
    )

    models = client.models.list()

    model = models.data[0].id

    llm = LangchainOpenAI(
        model=model,
        openai_api_key=nr_openai_api_key,
        openai_api_base=llm_openai_api_base,
    )

    ## Load and chunk contents of the blog
    loader = WebBaseLoader(
        web_paths=("https://lilianweng.github.io/posts/2023-06-23-agent/",),
        bs_kwargs=dict(
            parse_only=bs4.SoupStrainer(
                class_=("post-content", "post-title", "post-header")
            )
        ),
    )
    docs = loader.load()

    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=512, chunk_overlap=128
    )
    all_splits = text_splitter.split_documents(docs)

    ## Index chunks
    vector_store = InMemoryVectorStore(embeddings)
    document_ids = vector_store.add_documents(documents=all_splits)
    logger.debug(document_ids)

    ## Define prompt for question-answering
    prompt = hub.pull("rlm/rag-prompt")

    ## Define state for application
    class State(TypedDict):
        question: str
        context: List[Document]
        answer: str

    ## Define application steps
    def retrieve(state: State):
        retrieved_docs = vector_store.similarity_search(state["question"])
        return {"context": retrieved_docs}

    def generate(state: State):
        docs_content = "\n\n".join(doc.page_content for doc in state["context"])
        messages = prompt.invoke(
            {"question": state["question"], "context": docs_content}
        )
        response = llm.invoke(messages)
        return {"answer": response}

    ## Compile application and test
    graph_builder = StateGraph(State).add_sequence([retrieve, generate])
    graph_builder.add_edge(START, "retrieve")
    graph = graph_builder.compile()
    response = graph.invoke({"question": "What is Task Decomposition?"})
    logger.debug(response["answer"])
    logger.info("Finished LangChain RAG test")