Merge pull request #10 from LLMSQL/6-add-dependency-manager

pdm added; mypy and ruff added;

Author: DzmitryPihulski

.gitignore

@@ -5,3 +5,4 @@ dataset/sqlite_tables.db
 dist/
 
 *.egg-info/
+.pdm-python

llmsql/__init__.py

@@ -1,10 +1,7 @@
-import logging
-import os
-
 __version__ = "0.1.2"
 
 
-def __getattr__(name: str):
+def __getattr__(name: str):  # type: ignore
     if name == "LLMSQLVLLMInference":
         from .inference.inference import LLMSQLVLLMInference
 

llmsql/__main__.py

@@ -2,7 +2,7 @@
 import sys
 
 
-def main():
+def main() -> None:
     parser = argparse.ArgumentParser(prog="llmsql", description="LLMSQL CLI")
     subparsers = parser.add_subparsers(dest="command")
 

llmsql/evaluation/evaluator.py

@@ -22,9 +22,8 @@
 
 import json
 import os
-import sqlite3
 from pathlib import Path
-from typing import Dict, List, Optional
+import sqlite3
 
 from huggingface_hub import hf_hub_download
 from rich.progress import track
@@ -45,7 +44,7 @@ def __init__(self, workdir_path: str = "llmsql_workdir"):
         """
         Initialize evaluator.
         """
-        self.conn = None
+        self.conn: sqlite3.Connection | None = None
         self.workdir_path = Path(workdir_path)
         self.workdir_path.mkdir(parents=True, exist_ok=True)
         self.repo_id = "llmsql-bench/llmsql-benchmark"
@@ -67,30 +66,27 @@ def _download_file(self, filename: str) -> str:
         log.info(f"File saved at: {file_path}")
         return file_path
 
-    def connect(self, db_path: str):
+    def connect(self, db_path: str) -> None:
         """Establish SQLite connection."""
         if not os.path.exists(db_path):
             raise FileNotFoundError(f"Database not found at {db_path}")
         self.conn = sqlite3.connect(db_path)
 
-    def close(self):
+    def close(self) -> None:
         """Close SQLite connection if open."""
         if self.conn:
             self.conn.close()
             self.conn = None
 
-    # ------------------------------------------------------------------
-    # Evaluation
-    # ------------------------------------------------------------------
     def evaluate(
         self,
         outputs_path: str,
-        questions_path: Optional[str] = None,
-        db_path: Optional[str] = None,
-        save_report: Optional[str] = None,
+        questions_path: str | None = None,
+        db_path: str | None = None,
+        save_report: str | None = None,
         show_mismatches: bool = True,
         max_mismatches: int = 5,
-    ) -> Dict:
+    ) -> dict:
         """
         Evaluate predicted SQL queries against benchmark ground truth.
 
@@ -145,12 +141,14 @@ def evaluate(
             "gold_none": 0,
             "sql_errors": 0,
         }
-        mismatches: List[Dict] = []
+        mismatches: list[dict] = []
 
         for item in track(outputs, description="Evaluating"):
             metrics["total"] += 1
             is_match, mismatch_info, m_update = evaluate_sample(
-                item, questions, self.conn
+                item,
+                questions,
+                self.conn,  # type: ignore
             )
 
             metrics["matches"] += is_match

llmsql/finetune/finetune.py

@@ -20,23 +20,24 @@
 """
 
 import argparse
+from collections.abc import Callable
 import os
-import shutil
 from pathlib import Path
-from typing import Dict, Optional
+import shutil
+from typing import Any
 
-import torch
-import yaml
-from datasets import Dataset
+from datasets import Dataset  # type: ignore
 from huggingface_hub import hf_hub_download
+import torch
 from transformers import AutoModelForCausalLM
 from trl import SFTConfig, SFTTrainer
+import yaml
 
 from llmsql.loggers.logging_config import log
 from llmsql.utils.utils import choose_prompt_builder, load_jsonl
 
 
-def parse_args_and_config():
+def parse_args_and_config() -> argparse.Namespace:
     """Parse CLI args and optionally merge with YAML config."""
     p = argparse.ArgumentParser(
         description="Fine-tune a causal LM on Text-to-SQL benchmark."
@@ -66,10 +67,10 @@ def parse_args_and_config():
     # Load YAML config
     config = {}
     if args.get("config_file"):
-        with open(args["config_file"], "r") as f:
+        with open(args["config_file"]) as f:
             config = yaml.safe_load(f)
 
-    def flatten(d, parent_key="", sep="_"):
+    def flatten(d: Any, parent_key: str = "", sep: str = "_") -> dict[str, Any]:
         items = {}
         for k, v in d.items():
             new_key = f"{parent_key}{sep}{k}" if parent_key else k
@@ -88,7 +89,7 @@ def flatten(d, parent_key="", sep="_"):
     return argparse.Namespace(**args)
 
 
-def build_dataset(file_path: str, tables: Dict, prompt_builder) -> Dataset:
+def build_dataset(file_path: str, tables: dict, prompt_builder: Callable) -> Dataset:
     """Convert JSONL file to HF dataset samples."""
     questions = load_jsonl(file_path)
     samples = []
@@ -115,15 +116,13 @@ def _download_file(filename: str, repo_id: str, workdir_path: str) -> str:
     shutil.copy(cached_path, local_path)
     return local_path
 
-    return cached_path
-
 
 def main(
     model_name_or_path: str,
     output_dir: str,
-    train_file: Optional[str] = None,
-    val_file: Optional[str] = None,
-    tables_file: Optional[str] = None,
+    train_file: str | None = None,
+    val_file: str | None = None,
+    tables_file: str | None = None,
     shots: int = 5,
     num_train_epochs: int = 3,
     per_device_train_batch_size: int = 4,
@@ -136,11 +135,11 @@ def main(
     max_length: int = 32768,
     no_eval: bool = False,
     eval_steps: int = 100,
-    wandb_project: Optional[str] = None,
-    wandb_run_name: Optional[str] = None,
-    wandb_key: Optional[str] = None,
+    wandb_project: str | None = None,
+    wandb_run_name: str | None = None,
+    wandb_key: str | None = None,
     wandb_offline: bool = False,
-):
+) -> None:
     os.makedirs(output_dir, exist_ok=True)
 
     # Seed
@@ -268,7 +267,7 @@ def main(
     log.info(f"Model saved at {output_dir}/final_model")
 
 
-def run_cli():
+def run_cli() -> None:
     args = parse_args_and_config()
     main(
         train_file=args.train_file,

llmsql/inference/inference.py

@@ -33,14 +33,14 @@
 """
 
 import os
-import random
 from pathlib import Path
-from typing import Dict, List, Optional
+import random
+from typing import Any
 
-import numpy as np
-import torch
 from dotenv import load_dotenv
 from huggingface_hub import hf_hub_download
+import numpy as np
+import torch
 from tqdm import tqdm
 from vllm import LLM, SamplingParams
 
@@ -54,8 +54,8 @@
 
 load_dotenv()
 
-Question = Dict[str, object]
-Table = Dict[str, object]
+Question = dict[str, object]
+Table = dict[str, object]
 
 
 class LLMSQLVLLMInference:
@@ -70,7 +70,7 @@ def __init__(
         tensor_parallel_size: int = 1,
         seed: int = 42,
         workdir_path: str = "llmsql_workdir",
-        **llm_kwargs,
+        **llm_kwargs: Any,
     ):
         """
         Initialize vLLM model for SQL inference.
@@ -127,15 +127,15 @@ def _download_file(self, filename: str) -> str:
     def generate(
         self,
         output_file: str,
-        questions_path: Optional[str] = None,
-        tables_path: Optional[str] = None,
+        questions_path: str | None = None,
+        tables_path: str | None = None,
         shots: int = 5,
         batch_size: int = 8,
         max_new_tokens: int = 256,
         temperature: float = 1.0,
         do_sample: bool = True,
-        **sampling_kwargs,
-    ) -> List[Dict[str, str]]:
+        **sampling_kwargs: Any,
+    ) -> list[dict[str, str]]:
         """
         Generate SQL queries for all benchmark questions.
 
@@ -199,7 +199,7 @@ def generate(
         prompt_builder = choose_prompt_builder(shots)
         log.info(f"Using {shots}-shot prompt builder: {prompt_builder.__name__}")
 
-        all_results: List[Dict[str, str]] = []
+        all_results: list[dict[str, str]] = []
         total = len(questions)
 
         temperature = 0.0 if not do_sample else temperature
@@ -225,8 +225,8 @@ def generate(
             # Generate with vLLM
             outputs = self.llm.generate(prompts, sampling_params)
 
-            batch_results: List[Dict[str, str]] = []
-            for q, out in zip(batch, outputs):
+            batch_results: list[dict[str, str]] = []
+            for q, out in zip(batch, outputs, strict=False):
                 text = out.outputs[0].text
                 batch_results.append(
                     {

llmsql/prompts/prompts.py

@@ -1,4 +1,9 @@
-def build_prompt_5shot(question, headers, types, sample_row) -> str:
+def build_prompt_5shot(
+    question: str,
+    headers: list[str],
+    types: list[str],
+    sample_row: list[str | float | int],
+) -> str:
     return f"""You are an expert SQLite SQL query generator.
 Your task: Given a question and a table schema, output ONLY a valid SQL SELECT query.
 ⚠️ STRICT RULES:
@@ -52,7 +57,12 @@ def build_prompt_5shot(question, headers, types, sample_row) -> str:
 SQL:"""
 
 
-def build_prompt_1shot(question, headers, types, sample_row) -> str:
+def build_prompt_1shot(
+    question: str,
+    headers: list[str],
+    types: list[str],
+    sample_row: list[str | float | int],
+) -> str:
     return f"""You are an expert SQLite SQL query generator.
 Your task: Given a question and a table schema, output ONLY a valid SQL SELECT query.
 ⚠️ STRICT RULES:
@@ -78,7 +88,12 @@ def build_prompt_1shot(question, headers, types, sample_row) -> str:
 SQL:"""
 
 
-def build_prompt_0shot(question, headers, types, sample_row) -> str:
+def build_prompt_0shot(
+    question: str,
+    headers: list[str],
+    types: list[str],
+    sample_row: list[str | float | int],
+) -> str:
     return f"""You are an expert SQLite SQL query generator.
 Your task: Given a question and a table schema, output ONLY a valid SQL SELECT query.
 ⚠️ STRICT RULES:

llmsql/utils/evaluation_utils.py

@@ -1,11 +1,11 @@
 import sqlite3
-from typing import List, Optional, Tuple
+from typing import Any
 
 from llmsql.loggers.logging_config import log
 from llmsql.utils.regex_extractor import find_sql
 
 
-def execute_sql(conn: sqlite3.Connection, sql: str) -> Optional[List[Tuple]]:
+def execute_sql(conn: sqlite3.Connection, sql: str) -> list[tuple] | None:
     """
     Execute a SQL query on the given SQLite connection and return its results.
 
@@ -59,7 +59,11 @@ def fix_table_name(sql: str, table_id: str) -> str:
     )
 
 
-def evaluate_sample(item, questions, conn):
+def evaluate_sample(
+    item: dict[str, int | str],
+    questions: dict[int, dict[str, str]],
+    conn: sqlite3.Connection,
+) -> tuple[int, dict[str, Any] | None, dict[Any, Any]]:
     """
     Evaluate a single model prediction against the gold (ground-truth) SQL query.
 
@@ -93,6 +97,9 @@ def evaluate_sample(item, questions, conn):
     """
     # Extract question metadata
     qid = item["question_id"]
+    assert isinstance(qid, int), (
+        "question_id in the outputs file needs to be of type int."
+    )
     q_info = questions[qid]
     table_id, gold_sql, question_text = (
         q_info["table_id"],
@@ -115,6 +122,9 @@ def evaluate_sample(item, questions, conn):
     last_pred_res = None  # store last prediction results for mismatch logging
 
     # Loop over all SQL queries extracted from the model output
+    assert isinstance(item["completion"], str), (
+        f"Completion filed in outputs file must be of type string: {item['completion']}. Type: {type(item['completion'])}"
+    )
     for pred_sql in find_sql(item["completion"]):
         # Replace placeholder table names with the actual one
         pred_sql_fixed = fix_table_name(pred_sql, table_id)

llmsql/utils/regex_extractor.py

@@ -1,8 +1,7 @@
 import re
-from typing import List
 
 
-def find_sql(model_output: str, limit: int = 10) -> List[str]:
+def find_sql(model_output: str, limit: int = 10) -> list[str]:
     """Function to extract SQL queries from the model's response
 
     Args: