lion_linker

lion_linker is a Python library that uses Large Language Models (LLMs) to perform entity linking over tabular data. It efficiently links entity mentions in tables to relevant knowledge graph entities using customizable prompts and batch processing.

Overview

lion_linker simplifies the process of entity linking by leveraging LLMs to identify and connect textual data in tables to relevant entities in a knowledge base. It is highly customizable, making it ideal for projects involving research, NLP, and large-scale data processing.

Features

• Entity Linking with LLMs: Provides accurate and context-aware entity linking using large language models.
• Command Line Interface (CLI): Process CSV files via the command line for ease of use and automation.
• Customizable Prompt Templates: Adjust the linking prompt to fit your data structure.
• Scalable for Large Datasets: Process large datasets by batching the work and customizing API result limits.
• Flexible API: Programmatic interface for advanced users.
• REST API (FastAPI): Host LionLinker as a background job service with queued workloads.

Installation

To use lion_linker, you will need to install both the Python package and Ollama, which must be running locally.

Step 1: Install lion_linker via pip

You can install the latest version directly from the GitHub repository using the following command:

pip install git+https://github.com/enRichMyData/lion_linker.git

Alternatively, clone the repository and install from source:

git clone https://github.com/enRichMyData/lion_linker.git
cd lion_linker
pip install -e .

To use TableLlama run the following:

pip install git+https://github.com/enrichMyData/lion_linker.git#egg=lion_linker[tablellama]
pip install -U huggingface-hub

Step 2: Install Ollama

Ollama is required for the large language models to function with lion_linker. You must install it separately and ensure it is running locally.

1. Download and install Ollama from Ollama's official website.
2. After installation, start the Ollama service locally by running the following command:

ollama serve

This ensures that the LLM models can be accessed by lion_linker during the entity linking process.

Step 3:

lion_linker requires certain environment variables to be set, such as the API_URL and API_TOKEN for the entity retrieval API. To make this process easier, a .env.template file is provided in the repository.

1. Create a .env file by copying .env.template:

cp .env.template .env

2. Edit the .env file to add your specific API details. Open .env in a text editor and fill in the required values:

RETRIEVER_ENDPOINT=https://lamapi.hel.sintef.cloud/lookup/entity-retrieval
RETRIEVER_TOKEN=your_api_token  # Replace with your actual API token

The .env file will be used to securely store your Retriever credentials and other sensitive configuration data, so make sure it is not committed to version control. Moreover, if one wants to use OpenRouter the following additional credential must be set:

OPENAI_API_KEY=sk-v1-...

which can be retreived from the OpenRouter settings. If you plan to use Cerebras Cloud, set the following instead:

CEREBRAS_API_KEY=....

You can create or rotate Cerebras keys from the Cerebras console.

3. Verify the .env file by checking that RETRIEVER_ENDPOINT and RETRIEVER_TOKEN are correctly set, as these values will be automatically loaded by lion_linker when it runs.

Usage

Python Example

import os
import subprocess

from dotenv import load_dotenv

from lion_linker.lion_linker import LionLinker
from lion_linker.retrievers import LamapiClient

# Load environment variables from the .env file
load_dotenv()

# Define necessary file paths and parameters
input_csv = "tests/data/film.csv"
prompt_file_path = "lion_linker/prompt/prompt_template.txt"
model_name = "gemma2:2b"  # Use the correct model name
output_csv = "output_test.csv"
chunk_size = 16  # How many rows to process
num_candidates = 20  # Maximum number of candidates from the Retriever per mention
format_candidates = False  # Use TableLlama-style formatting instead of JSON
compact_candidates = False  # Use compact non-JSON list when format_candidates is False
max_parallel_prompts = 1  # Maximum number of prompt batches to send concurrently

# Load API parameters from environment variables
retriever_endpoint = os.getenv("RETRIEVER_ENDPOINT")
retriever_token = os.getenv("RETRIEVER_TOKEN")

# Additional parameters as per the latest LionLinker version
mention_columns = ["title"]  # Columns to link entities from
model_api_provider = "ollama"  # e.g. "ollama", "openrouter", "huggingface", or "cerebras"
ollama_host = "http://localhost:11434"  # Default Ollama host if not specified it will use the Default Ollama host anyway
model_api_key = None  # Optional model API key if required
gt_columns = []  # Specify any ground truth columns to exclude for testing

# Initialize the retriever instance
retriever = LamapiClient(retriever_endpoint, retriever_token, num_candidates=num_candidates)

# Initialize the LionLinker instance
lion_linker = LionLinker(
    input_csv=input_csv,
    model_name=model_name,
    retriever=retriever,
    output_csv=output_csv,
    prompt_file_path=prompt_file_path,
    chunk_size=chunk_size,
    mention_columns=mention_columns,
    compact_candidates=compact_candidates,
    model_api_provider=model_api_provider,
    ollama_host=ollama_host,
    model_api_key=model_api_key,
    gt_columns=gt_columns,
    format_candidates=format_candidates,
    max_parallel_prompts=max_parallel_prompts,
)

# Start the Ollama server as a background process
process = subprocess.Popen(["ollama", "serve"])

# Run the entity linking
await lion_linker.run()

# Stop the Ollama server
process.terminate()

Using Cerebras Cloud

To call hosted models via Cerebras Cloud, set model_api_provider="cerebras" and supply a Cerebras model name (for example llama-3.3-70b). You can pass the key explicitly or rely on the CEREBRAS_API_KEY environment variable.

import os

from lion_linker.lion_linker import LionLinker

lion_linker = LionLinker(
    input_csv=input_csv,
    model_name="llama-3.3-70b",
    retriever=retriever,
    mention_columns=mention_columns,
    model_api_provider="cerebras",
    model_api_key=os.getenv("CEREBRAS_API_KEY"),
)

CLI Example

python -m lion_linker.cli \
  --lion.input_csv "./data/film.csv" \
  --lion.model_name "gemma2:2b" \
  --lion.mention_columns '[title]' \
  --lion.ollama_host "http://localhost:11434" \
  --retriever.class_path lion_linker.retrievers.LamapiClient \
  --retriever.endpoint "https://lamapi.hel.sintef.cloud/lookup/entity-retrieval" \
  --retriever.token "lamapi_demo_2023" \
  --retriever.kg wikidata \
  --retriever.num_candidates 5 \
  --retriever.cache False

If one wants to change the retriever and for example use the Wikidata Lookup Service instead, the following can be used instead:

python -m lion_linker.cli \
  --lion.input_csv "./data/film.csv" \
  --lion.model_name "gemma2:2b" \
  --lion.mention_columns '[title]' \
  --lion.ollama_host "http://localhost:11434" \
  --retriever.class_path lion_linker.retrievers.WikidataClient \
  --retriever.endpoint "https://query.wikidata.org/sparql" \
  --retriever.language "en" \
  --retriever.num_candidates 5

Another possibility is to retrieve candidates for mentions through OpenRefine:

python -m lion_linker.cli \
  --lion.input_csv "./data/film.csv" \
  --lion.model_name "gemma2:2b" \
  --lion.mention_columns '[title]' \
  --lion.ollama_host "http://localhost:11434" \
  --retriever.class_path lion_linker.retrievers.OpenRefineClient \
  --retriever.endpoint "https://wikidata.reconci.link/api" \
  --retriever.num_candidates 5

Explanation of Parameters

• input_csv: Path to your input CSV file.

• output_csv: Path where the output file will be saved.

• ollama_host: The host where the Ollama service is running.

• --prompt_file_path: Path to a file containing a custom prompt template.

• --model: The LLM model to use for entity linking.

• --chunk_size: Defines how many rows to process at once.

• --mention_columns: Columns in the CSV that contain entity mentions.

• --num_candidates: Maximum number of candidates returned by the API per mention.

• --format_candidates: Use TableLlama/legacy bracket formatting instead of JSON candidates.

• --compact_candidates: Use the compact list format (ID | TYPE | DESCRIPTION) when JSON is disabled.

• --max_parallel_prompts: Maximum number of prompt batches to send concurrently.

REST API server

The repository includes a FastAPI application (located in app/) that exposes a small REST interface for running LionLinker jobs as queued compute tasks. MongoDB is the job store and the queue, and prediction rows are written to a separate predictions collection.

Install dependencies

pip install -e .[app]

Configure

The API shares the same environment variables as the Python/CLI interfaces. The most relevant flags are:

• RETRIEVER_ENDPOINT, RETRIEVER_TOKEN: LamAPI (or compatible) endpoint configuration.
• LionLinker-specific settings (model, prompts, mention columns, etc.) should be supplied per request via config.lion / config.retriever in the job payload.
• LION_API_KEY: required API key; send it via X-API-Key or Authorization: Bearer.
• LION_JOB_SECRET_TTL_SECONDS: optional TTL for per-job model API keys (defaults to 24h).
• LION_DRY_RUN=true: force offline/dry runs that emit ANSWER:NIL predictions without contacting retrievers or models (handy for local testing).
• MONGO_URI (defaults to mongodb://localhost:27017), MONGO_DB, and MONGO_COLLECTION_PREFIX: connection settings for the MongoDB instance backing the job store.
• WORKSPACE_PATH (defaults to data/api_runs): where intermediate CSV/JSON artifacts are written inside the container.
• PREDICTION_BATCH_ROWS: unused in the queue API (left for compatibility with older data formats).

Per-job model keys can be supplied as config.lion.model_api_key in the job payload or via the X-LLM-API-Key header. The API stores them in a short-lived secrets collection and does not persist them in the job config.

Run

uvicorn app.main:app --reload

Detailed API documentation is available in docs/api_reference.md and via the autogenerated Swagger UI at /docs.

Docker Compose

You can spin up the API and a MongoDB instance together using the provided compose file:

cd docker/service
docker compose up --build

The compose stack exposes the API on http://localhost:9000, MongoDB on mongodb://localhost:27017, and mounts the host data/ directory into the container at /app/data so persisted predictions (CSV, JSON, MongoDB volume) are easy to inspect.

Endpoints

• GET /health – liveness check.
• GET /capabilities – supported modes, limits, defaults.
• POST /jobs – create a job and return job_id.
• GET /jobs/{job_id} – job status + progress.
• GET /jobs/{job_id}/results – cursor-paged final decisions (no candidates).
• GET /jobs/{job_id}/cells/{row}/{col}/candidates – candidates for one cell.
• POST /jobs/{job_id}:cancel – cancel a queued job.

Job request (inline example)

{
  "table_id": "my_table_001",
  "input": {
    "mode": "inline",
    "format": "application/json",
    "table": {
      "header": ["Point of Interest", "Place"],
      "rows": [
        {"row_id": "r1", "cells": ["John F. Kennedy Presidential Library and Museum", "Columbia Point"]},
        {"row_id": "r2", "cells": ["Petrie Museum of Egyptian Archaeology", "London"]}
      ]
    }
  },
  "link_columns": ["Point of Interest"],
  "top_k": 5,
  "execution": "async"
}

Example curl flow

curl -s -X POST "http://localhost:9000/jobs" \
  -H "Content-Type: application/json" \
  -H "X-API-Key: ${LION_API_KEY}" \
  -d '{
    "table_id": "demo_table",
    "input": {
      "mode": "inline",
      "format": "application/json",
      "table": {
        "header": ["Title", "Place"],
        "rows": [
          {"row_id": "r1", "cells": ["The Matrix", "USA"]},
          {"row_id": "r2", "cells": ["Spirited Away", "Japan"]}
        ]
      }
    },
    "link_columns": ["Title"],
    "top_k": 5,
    "execution": "async"
  }'

curl -s "http://localhost:9000/jobs/<job_id>" \
  -H "X-API-Key: ${LION_API_KEY}"

curl -s "http://localhost:9000/jobs/<job_id>/results" \
  -H "X-API-Key: ${LION_API_KEY}"

Migration notes

• The queue API uses {prefix}_jobs and {prefix}_predictions (or jobs/predictions if the prefix is empty).
• Legacy lion_* dataset/table collections are no longer used by the API.
• Old jobs with statuses like pending/completed will not be claimed by the new worker; archive or delete them if needed.

Running Tests

You can run the tests with:

python -m unittest discover -s tests

This will execute all unit tests in the tests/ directory.

License

This project is licensed under the terms of the Apache License 2.0.