Skip to content

Instantly share code, notes, and snippets.

@wware
Last active February 24, 2025 21:51
Show Gist options
  • Save wware/561ebf2d8fa65ae81797dfc90569db3e to your computer and use it in GitHub Desktop.
Save wware/561ebf2d8fa65ae81797dfc90569db3e to your computer and use it in GitHub Desktop.
docker-compose.yml
Dockerfile*
models/
ollama_data/
.pytest_cache/
__pycache__/
LANGUAGE_MODEL=llama2
# LANGUAGE_MODEL=smollm2:135m
# LANGUAGE_MODEL=qwen2.5-coder:1.5b
# LANGUAGE_MODEL=codellama
OLLAMA_HOST=http://ollama:11434
# where models are stored on the host machine
OLLAMA_MODELS=./models
models/
ollama_data/
.pytest_cache/
__pycache__/

Jupyter AI with Ollama

https://bit.ly/jupyter-ai-ollama

This project combines Jupyter AI with Ollama to create a local, containerized environment for AI-assisted data analysis and development. It uses Docker Compose to manage two services:

  1. An Ollama service for running large language models locally
  2. A Jupyter service with Jupyter AI integration

Prerequisites

  • Docker and Docker Compose
  • Sufficient disk space for language models (varies by model)
  • Optional: NVIDIA container toolkit (for GPU support)

Quick Start

  1. Clone this repository
  2. Configure your model in .env (default is qwen2.5-coder:1.5b)
  3. Start the services:
docker-compose up -d
  1. Pull the language model manually (recommended):
docker-compose exec -it ollama bash
ollama pull $LANGUAGE_MODEL  # as specified in .env

While the code includes automatic model pulling, manual pulling is recommended to monitor the download progress and handle any issues.

  1. Access Jupyter Lab at http://localhost:8888 (token: 123)

Configuration

Environment Variables

Edit .env to configure:

  • LANGUAGE_MODEL: Choose your Ollama model
  • OLLAMA_HOST: Internal service URL (default: http://ollama:11434)

Available Models

These are some recommended models to get started:

Browse the Ollama Model Library for many more available models. You can use any model from the library by updating the LANGUAGE_MODEL in your .env file.

Project Structure

.
├── .dockerignore       # Excludes unnecessary files from builds
├── .env               # Environment configuration
├── Dockerfile         # Jupyter service configuration
├── Dockerfile.ollama  # Ollama service configuration
├── config.json       # Jupyter AI configuration
├── docker-compose.yml # Service orchestration
├── entrypoint.sh    # Ollama initialization
└── jupyter-start.sh # Jupyter initialization

GPU Support (Optional)

This project runs well on CPU-only systems. However, if you have a compatible NVIDIA GPU, you can enable GPU acceleration:

  1. Install the NVIDIA Container Toolkit:
curl -fsSL https://nvidia.github.io/libnvidia-container/gpgkey | sudo gpg --dearmor -o /usr/share/keyrings/nvidia-container-toolkit-keyring.gpg
curl -s -L https://nvidia.github.io/libnvidia-container/stable/deb/nvidia-container-toolkit.list | \
  sed 's#deb https://#deb [signed-by=/usr/share/keyrings/nvidia-container-toolkit-keyring.gpg] https://#g' | \
  sudo tee /etc/apt/sources.list.d/nvidia-container-toolkit.list
sudo apt update
sudo apt install -y nvidia-container-toolkit
  1. Add GPU configuration to docker-compose.yml under the ollama service:
deploy:
  resources:
    reservations:
      devices:
        - driver: nvidia
          count: 1
          capabilities: [gpu]

Resource Usage

Resource requirements vary by model:

  • Small models (smollm2:135m): 2-4GB RAM
  • Medium models (qwen2.5-coder:1.5b): 4-8GB RAM
  • Large models (llama2): 8GB+ RAM

Troubleshooting

If the model fails to load:

  1. Stop the services: docker-compose down
  2. Remove the ollama_data volume
  3. Restart and manually pull the model as described above

For more issues, check container logs:

docker-compose logs -f jupyter
docker-compose logs -f ollama

Perplexity.AI suggests improvements:

  1. Automated Model Download: The project suggests manually pulling the language model, though automatic model pulling is included in the code. Improving the reliability and transparency of the automatic pulling process would enhance the user experience. Consider adding progress bars or more detailed logging during the download process. Error handling could also be improved to automatically retry downloads or suggest alternative models if the specified model is unavailable.

  2. Simplified GPU Configuration: While GPU support is mentioned, enabling it requires manual configuration of the docker-compose.yml file. Providing a simpler way to enable GPU support (e.g., through an environment variable or a separate docker-compose.gpu.yml file) would make it more accessible to users with NVIDIA GPUs.

  3. Health Checks: Implement Docker health checks for both the Ollama and Jupyter services. This would allow Docker Compose to automatically restart services if they become unhealthy, improving the overall robustness of the environment.

  4. Dynamic Resource Allocation: Explore options for dynamic resource allocation based on the selected language model. This could involve adjusting the memory and CPU resources allocated to the Ollama service based on the model's requirements, optimizing resource utilization.

  5. GUI for Model Selection: Consider adding a simple GUI or CLI tool to select and manage language models. This would make it easier for users to discover and switch between different models without manually editing the .env file.

  6. Security Considerations: While the project uses a default token for JupyterLab, it's essential to emphasize the importance of changing this token for security reasons. Consider adding a warning message or instructions on how to generate a secure token during the setup process.

  7. Expand Documentation: Add more detailed documentation on how to use Jupyter AI with Ollama, including examples of common use cases and troubleshooting tips. This would help users get started quickly and effectively.

  8. Pre-configured notebooks: Include sample Jupyter notebooks that demonstrate how to use the integrated environment for common tasks, like data analysis, code generation, and model exploration. This could serve as a quick start guide for new users.

  9. Model management: Provide clear instructions on how to manage different language models, including how to remove models that are no longer needed and how to update models to the latest versions.

  10. Multi-platform support: Expand the documentation to include specific instructions for setting up the environment on different operating systems (e.g., Windows, macOS, Linux).

{
"model_provider_id": "ollama:LANGUAGE_MODEL",
"embeddings_provider_id": "ollama:nomic-embed-text",
"send_with_shift_enter": false,
"fields": {},
"api_keys": {},
"completions_model_provider_id": null,
"completions_fields": {}
}
version: '3'
services:
ollama:
build:
context: .
dockerfile: Dockerfile.ollama
ports:
- "11434:11434"
environment:
- LANGUAGE_MODEL
volumes:
- ./ollama_data:/root/.ollama
- ${OLLAMA_MODELS:-./models}:/models
jupyter:
build:
context: .
dockerfile: Dockerfile
ports:
- "8888:8888"
environment:
- LANGUAGE_MODEL
- OLLAMA_HOST
depends_on:
- ollama
volumes:
ollama_data:
FROM jupyter/scipy-notebook:latest
RUN pip install 'jupyter-ai[all]' pytest ipytest
COPY config.json /tmp/config.json
ENV JUPYTER_ENABLE_LAB=yes
ENV JUPYTER_TOKEN=123
USER root
COPY jupyter-start.sh /jupyter-start.sh
RUN chmod +x /jupyter-start.sh
COPY hack.ipynb /home/jovyan/hack.ipynb
RUN chown jovyan:users /home/jovyan/hack.ipynb
COPY testit.sh /home/jovyan/testit.sh
RUN chown jovyan:users /home/jovyan/testit.sh
USER jovyan
ENTRYPOINT ["/jupyter-start.sh"]
FROM ollama/ollama:latest
COPY entrypoint.sh /entrypoint.sh
RUN chmod +x /entrypoint.sh
ENTRYPOINT ["/entrypoint.sh"]
#!/bin/sh
ollama serve &
sleep 5
ollama pull nomic-embed-text
ollama pull ${LANGUAGE_MODEL}
wait
Display the source blob
Display the rendered blob
Raw
{
"cells": [
{
"cell_type": "markdown",
"id": "96369004-9b12-4605-9160-56d6ac004baf",
"metadata": {},
"source": [
"# Here is some stuff\n",
"\n",
"Write a factorial function."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "d522fd9c-8817-4fcf-bccf-ee3b5d12dbe1",
"metadata": {},
"outputs": [],
"source": [
"def fac_helper(prod, n):\n",
" if not isinstance(n, int):\n",
" raise TypeError(n)\n",
" if n < 0:\n",
" raise ValueError(n)\n",
" if n < 2:\n",
" return prod\n",
" return fac_helper(prod * n, n - 1)\n",
"\n",
"def fac(n):\n",
" return fac_helper(1, n)"
]
},
{
"cell_type": "markdown",
"id": "8f697f8e-091b-44e2-a0a1-af27c3121165",
"metadata": {},
"source": [
"Next, write some pytest test cases for it."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "27e6e20a-8f44-406c-bdc0-fcc0c00be5c3",
"metadata": {},
"outputs": [],
"source": [
"import pytest\n",
"\n",
"@pytest.mark.parametrize(\"n, expected\", [\n",
" (0, 1), # 0! = 1\n",
" (1, 1), # 1! = 1\n",
" (5, 120), # 5! = 120\n",
" (10, 3628800) # 10! = 3,628,800\n",
"])\n",
"def test_fac_valid(n, expected):\n",
" assert fac(n) == expected\n",
"\n",
"@pytest.mark.parametrize(\"n\", [-1, -5])\n",
"def test_fac_invalid_negative(n):\n",
" with pytest.raises(ValueError):\n",
" fac(n)\n",
"\n",
"@pytest.mark.parametrize(\"n\", [2.5, \"3\", []])\n",
"def test_fac_invalid_type(n):\n",
" with pytest.raises(TypeError):\n",
" fac(n)"
]
},
{
"cell_type": "markdown",
"id": "96369004-9b12-4605-9160-56d6ac004bae",
"metadata": {},
"source": [
"# Run pytest"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "7e5fb5af-8e3a-49e0-b86b-8bff3a27a454",
"metadata": {},
"outputs": [],
"source": [
"import ipytest\n",
"ipytest.run()"
]
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3 (ipykernel)",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.11.6"
}
},
"nbformat": 4,
"nbformat_minor": 5
}
#!/bin/bash -xe
mkdir -p /home/jovyan/.local/share/jupyter/jupyter_ai
sed "s/ollama:LANGUAGE_MODEL/ollama:${LANGUAGE_MODEL}/" \
< /tmp/config.json \
> /home/jovyan/.local/share/jupyter/jupyter_ai/config.json
jupyter lab --ip=0.0.0.0 --allow-root
#!/bin/bash
TARGET="${1%.ipynb}.py"
python -m nbconvert --to script $1
python -m pytest $TARGET
rm -f $TARGET
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment