AI Engineer

5 minute read

Integrating Ollama with LangChain: A Comprehensive Guide to Self-Hosted Language Models

In the rapidly evolving landscape of AI development, the ability to run language models locally has become increasingly valuable. Local LLMs offer advantages in privacy, cost, and latency that cloud-based alternatives often can’t match. Ollama has emerged as a popular solution for running open-source language models locally, and LangChain provides a powerful framework for building applications with these models. In this guide, we’ll explore how to implement and optimize local language models using the OllamaLLM class in LangChain.

Why Use Local LLMs with Ollama?

Before diving into implementation details, let’s consider why you might want to use locally-hosted models:

  1. Privacy: Your data never leaves your machine
  2. Cost: No usage-based API charges
  3. Latency: Reduced network delays
  4. Offline capability: No internet connection required
  5. Customizability: Fine-tune and customize models for specific use cases

Getting Started with OllamaLLM

The OllamaLLM class in LangChain provides a convenient interface for interacting with Ollama models. It implements LangChain’s standard Runnable Interface, giving you access to all the standard LangChain functionality.

Installation

First, make sure you have Ollama installed on your system. You can find installation instructions at ollama.ai.

Next, install the LangChain Ollama integration:

pip install langchain-ollama

Basic Usage

Here’s a simple example of how to use OllamaLLM:

from langchain_ollama import OllamaLLM

# Initialize the LLM
llm = OllamaLLM(model="llama2")

# Generate a response
response = llm.invoke("Explain quantum computing in simple terms.")
print(response)

This code initializes an OllamaLLM instance using the “llama2” model and generates a response to a prompt about quantum computing.

Configuration Options

One of the strengths of OllamaLLM is its extensive configurability. Let’s explore some of the key parameters:

Model Selection and Base URL

llm = OllamaLLM(
    model="llama2",  # Choose your model
    base_url="http://localhost:11434"  # Default Ollama URL
)

Generation Parameters

You can fine-tune the generation behavior with parameters like:

llm = OllamaLLM(
    model="llama2",
    temperature=0.7,  # Controls creativity (default: 0.8)
    top_k=40,  # Limits token selection to top K options (default: 40)
    top_p=0.9,  # Nucleus sampling parameter (default: 0.9)
    num_predict=256,  # Maximum tokens to generate (default: 128)
)

Advanced Parameters

For more nuanced control, you can use parameters like:

llm = OllamaLLM(
    model="llama2",
    mirostat=1,  # Enable Mirostat sampling (0=disabled, 1=Mirostat, 2=Mirostat 2.0)
    mirostat_tau=5.0,  # Controls coherence vs. diversity balance
    mirostat_eta=0.1,  # Response rate to feedback
    repeat_penalty=1.1,  # Penalize repetition (default: 1.1)
    seed=42,  # Set seed for reproducibility
)

Handling Context and Memory

The context window size is a critical parameter for LLMs. With OllamaLLM, you can control this using the num_ctx parameter:

llm = OllamaLLM(
    model="llama2",
    num_ctx=4096,  # Size of context window (default: 2048)
)

You can also control how the model handles repetition with parameters like:

llm = OllamaLLM(
    model="llama2",
    repeat_last_n=64,  # How far back to look to prevent repetition
    repeat_penalty=1.2,  # Higher values penalize repetition more strongly
)

Performance Optimization

To optimize performance on your hardware, you can adjust:

llm = OllamaLLM(
    model="llama2",
    num_gpu=1,  # Number of GPUs to use (default: 1 on macOS)
    num_thread=8,  # Number of CPU threads (defaults to optimal for your system)
)

Streaming Responses

LangChain makes it easy to stream responses from Ollama models:

from langchain_ollama import OllamaLLM

llm = OllamaLLM(model="llama2")

# Stream the response
for chunk in llm.stream("Write a short poem about programming."):
    print(chunk, end="", flush=True)

This is particularly useful for creating more responsive user interfaces.

Structured Output with JSON Mode

If you need structured data from your model, you can use the format parameter:

llm = OllamaLLM(
    model="llama2",
    format="json"  # Request JSON-formatted output
)

response = llm.invoke("List three programming languages and their key features")

Caching Responses

To improve performance and reduce redundant computation, you can enable caching:

from langchain.globals import set_llm_cache
from langchain.cache import InMemoryCache

# Set up a global cache
set_llm_cache(InMemoryCache())

# Enable caching in the LLM
llm = OllamaLLM(
    model="llama2",
    cache=True  # Use the global cache
)

Advanced Usage: Batch Processing

For processing multiple prompts efficiently:

prompts = [
    "What is machine learning?",
    "Explain natural language processing.",
    "How do neural networks work?"
]

# Process all prompts in parallel
responses = llm.batch(prompts)

for prompt, response in zip(prompts, responses):
    print(f"Prompt: {prompt}")
    print(f"Response: {response}")
    print("-" * 50)

Integration with LangChain’s Ecosystem

Since OllamaLLM implements LangChain’s Runnable interface, it integrates seamlessly with other LangChain components:

from langchain.prompts import PromptTemplate
from langchain.chains import LLMChain

prompt = PromptTemplate.from_template(
    "Explain {concept} in terms a {audience} would understand."
)

chain = LLMChain(llm=llm, prompt=prompt)

response = chain.invoke({
    "concept": "quantum entanglement", 
    "audience": "high school student"
})

print(response["text"])

Error Handling and Fallbacks

You can implement fallback mechanisms for robustness:

from langchain.schema.runnable import RunnableWithFallbacks

# Create a fallback LLM with different parameters
fallback_llm = OllamaLLM(
    model="llama2",
    temperature=0.2,  # More conservative settings
    num_predict=64
)

# Create a runnable with fallback
robust_llm = llm.with_fallbacks([fallback_llm])

# This will try the primary LLM first, then fall back if there's an error
response = robust_llm.invoke("Summarize the theory of relativity")

Saving and Loading Models

You can save your configured model for later use:

from pathlib import Path

# Save the configured LLM
llm.save("my_ollama_config.yaml")

# Later, you can load it
from langchain.llms import load_llm
loaded_llm = load_llm(Path("my_ollama_config.yaml"))

Token Counting and Context Management

Understanding token usage is important for managing context windows:

text = "This is a sample text to count tokens for."
token_count = llm.get_num_tokens(text)
print(f"Token count: {token_count}")

Conclusion

The OllamaLLM class in LangChain provides a powerful and flexible interface for working with locally-hosted language models. By leveraging Ollama’s capabilities through LangChain, developers can build sophisticated applications that benefit from the privacy, cost, and latency advantages of local LLMs while maintaining the flexibility and power of the LangChain ecosystem.

Whether you’re building a chatbot, a document analysis tool, or a creative writing assistant, the combination of Ollama and LangChain offers a compelling alternative to cloud-based LLM services. By understanding and optimizing the various configuration options available in OllamaLLM, you can tailor the behavior of your language model to meet your specific requirements and constraints.

Remember that local model performance will depend significantly on your hardware capabilities, so experiment with different models and settings to find the optimal balance for your use case.

This post was originally written in my native language and then translated using an LLM. I apologize if there are any grammatical inconsistencies.

You may also enjoy