Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
44 changes: 44 additions & 0 deletions .github/workflows/tests.yml
Original file line number Diff line number Diff line change
@@ -0,0 +1,44 @@
name: Tests

on:
push:
branches: [ main ]
pull_request:

jobs:
test:
runs-on: ubuntu-latest

steps:
- name: Checkout code
uses: actions/checkout@v4

- name: Set up Python
uses: actions/setup-python@v5
with:
python-version: "3.11"

- name: Install PDM
run: pip install pdm

- name: Cache PDM packages
uses: actions/cache@v4
with:
path: ~/.cache/pdm
key: ${{ runner.os }}-pdm-${{ hashFiles('pdm.lock') }}
restore-keys: |
${{ runner.os }}-pdm-

- name: Install dependencies (with dev)
run: pdm install --with dev

- name: Run tests with coverage
run: PYTHONPATH=. pdm run pytest --cov=llmsql --cov-report=xml --maxfail=1 --disable-warnings -v

- name: Upload coverage to Codecov
uses: codecov/codecov-action@v4
with:
token: ${{ secrets.CODECOV_TOKEN }}
files: ./coverage.xml
flags: unittests
fail_ci_if_error: true
3 changes: 2 additions & 1 deletion .gitignore
Original file line number Diff line number Diff line change
Expand Up @@ -5,4 +5,5 @@ dataset/sqlite_tables.db
dist/

*.egg-info/
.pdm-python
.pdm-python
.vscode
3 changes: 3 additions & 0 deletions README.md
Original file line number Diff line number Diff line change
@@ -1,3 +1,6 @@
[![codecov](https://codecov.io/gh/LLMSQL/llmsql-benchmark/branch/main/graph/badge.svg)](https://codecov.io/gh/LLMSQL/llmsql-benchmark)


# LLMSQL

Patched and improved version of the original large crowd-sourced dataset for developing natural language interfaces for relational databases, [WikiSQL](https://github.com/salesforce/WikiSQL).
Expand Down
4 changes: 4 additions & 0 deletions llmsql/inference/inference.py
Original file line number Diff line number Diff line change
Expand Up @@ -96,6 +96,10 @@ def __init__(
self.workdir_path.mkdir(parents=True, exist_ok=True)
self.repo_id = "llmsql-bench/llmsql-benchmark"

if "device" not in llm_kwargs:
llm_kwargs["device"] = "cuda" if torch.cuda.is_available() else "cpu"
self.device = llm_kwargs["device"]

log.info(
f"Loading vLLM model {model_name} with tensor_parallel_size={tensor_parallel_size}..."
)
Expand Down
30 changes: 29 additions & 1 deletion pdm.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

1 change: 1 addition & 0 deletions pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -57,6 +57,7 @@ dev = [
"pre-commit>=4.3.0",
"types-tqdm>=4.67.0.20250809",
"types-PyYAML>=6.0.12.20250915",
"pytest-asyncio>=1.2.0",
]

[tool.pdm]
Expand Down
176 changes: 0 additions & 176 deletions requirements.txt

This file was deleted.

70 changes: 70 additions & 0 deletions tests/conftest.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,70 @@
import json
import os
import sqlite3

import pytest

import llmsql.inference.inference as inference


@pytest.fixture
def temp_dir(tmp_path):
return tmp_path


@pytest.fixture
def dummy_db_file(tmp_path):
"""Create a temporary SQLite DB file for testing, cleanup afterwards."""
db_path = tmp_path / "test.db"
conn = sqlite3.connect(db_path)
conn.execute("CREATE TABLE test (id INTEGER PRIMARY KEY, name TEXT)")
conn.execute("INSERT INTO test (name) VALUES ('Alice'), ('Bob')")
conn.commit()
conn.close()

yield str(db_path)

# cleanup
if os.path.exists(db_path):
os.remove(db_path)


@pytest.fixture
def mock_llm(monkeypatch):
"""Mock vLLM LLM to avoid GPU/model loading."""

class DummyOutput:
def __init__(self, text="SELECT 1"):
self.outputs = [type("Obj", (), {"text": text})()]

class DummyLLM:
def generate(self, prompts, sampling_params):
return [DummyOutput(f"-- SQL for: {p}") for p in prompts]

monkeypatch.setattr(inference, "LLM", lambda **_: DummyLLM())
return DummyLLM()


@pytest.fixture
def fake_jsonl_files(tmp_path):
"""Create fake questions.jsonl and tables.jsonl."""
qpath = tmp_path / "questions.jsonl"
tpath = tmp_path / "tables.jsonl"

questions = [
{"question_id": "q1", "question": "How many users?", "table_id": "t1"},
{"question_id": "q2", "question": "List names", "table_id": "t1"},
]
tables = [
{
"table_id": "t1",
"header": ["id", "name"],
"types": ["int", "text"],
"rows": [[1, "Alice"], [2, "Bob"]],
}
]

qpath.write_text("\n".join(json.dumps(q) for q in questions))
tpath.write_text("\n".join(json.dumps(t) for t in tables))

return str(qpath), str(tpath)
Loading