Text-to-SQL with Memory Experiment

import os
import lamini
import pathlib
import datetime

from lamini.experiment.generators import (
    SchemaToSQLGenerator,
    SQLDebuggerGenerator,
    PatternQuestionGenerator,
    VariationQuestionGenerator,
    QuestionDecomposerGenerator
)
from lamini.experiment.validators import (
    SQLValidator
)

from helpers import (
    load_config,
    process_jsonl,
    read_jsonl, 
    get_schema, 
    format_glossary, 
    save_results, 
    generate_variations,
    process_variation
)

# Load configuration
config = load_config("config.yml")

# Set up Lamini API credentials
lamini.api_url = config['api']['url']
lamini.api_key = config['api']['key']

# Initialize generators and validators
generators = {
    "pattern": PatternQuestionGenerator(model=config['model']['default']),
    "variation": VariationQuestionGenerator(model=config['model']['default']),
    "decomposer": QuestionDecomposerGenerator(model=config['model']['default'])
}

sql_gen = SchemaToSQLGenerator(
    model=config['model']['default'],
    db_type="sqlite",
    db_params=str(config['database']['path']),
    schema=get_schema(config['database']['path'])
)

# Process evaluation questions to generate training data
# (See generate_data.py in the repository for the complete implementation)

import os
import json
import lamini
from typing import List, Dict
from lamini import ErrorAnalysis, SQLErrorAnalysis
from lamini.experiment.generators import (
    SchemaToSQLGenerator,
    SQLDebuggerGenerator
)
from lamini.experiment.validators import (
    SQLValidator
)

# Initialize ErrorAnalysis
error_analysis = ErrorAnalysis(
    model=config['model']['default'],
    schema=schema,
    glossary=formatted_glossary
)

# Analyze topic coverage
coverage_analysis = error_analysis.analyze_topic_coverage(
    gold_questions=gold_questions,
    generated_questions=generated_questions
)

# Generate additional questions for missing topics
additional_questions = []
if coverage_analysis["missing_topics"] or coverage_analysis["underrepresented_topics"]:
    additional_questions = error_analysis.generate_additional_questions(
        coverage_analysis=coverage_analysis,
        num_questions_per_topic=2
    )

# Analyze SQL errors
sql_error_analysis = SQLErrorAnalysis(model=config['model']['default'])
failed_queries = sql_error_analysis.extract_failed_queries(results_data)

# (See analyze_generated_data.py in the repository for the complete implementation)

import os
import lamini
from lamini import Lamini
from helpers import read_jsonl, load_config

# Load configuration
config = load_config("config.yml")

# Set up Lamini API credentials
lamini.api_url = config['api']['url']
lamini.api_key = config['api']['key']

# Read training data
input_file = "flattened_training_data.jsonl"
rows = read_jsonl(input_file)

# Configure Memory Tuning
model_name = config['model']['memory_tuning']
max_steps = config['memory_tuning']['max_steps']
learning_rate = config['memory_tuning']['learning_rate']
max_gpus = config['memory_tuning']['max_gpus']
max_nodes = config['memory_tuning']['max_nodes']

# Submit Memory Tuning job
llm = Lamini(model_name=model_name)
results = llm.tune(
    data_or_dataset_id=rows,
    finetune_args={
        "max_steps": max_steps,
        "learning_rate": learning_rate,
    },
    gpu_config={
        "max_gpus": max_gpus,
        "max_nodes": max_nodes
    }
)

# (See memory_tuning.py in the repository for the complete implementation)

import json
import os
import lamini
from lamini import Lamini
from helpers import save_results_to_jsonl, load_config
from lamini.experiment.error_analysis_eval import SQLExecutionPipeline

# Load configuration
config = load_config("config.yml")

# Set up Lamini API credentials
lamini.api_url = config['api']['url']
lamini.api_key = config['api']['key']

# Initialize the model with your fine-tuned model ID
model_id = "YOUR_TUNED_MODEL_ID"  # Replace with your model ID from Memory Tuning
llm = Lamini(model_name=model_id)

# Read test cases
test_cases = read_test_set(config['paths']['gold_test_set'])

# Run inference
inference_results = run_inference(test_cases, model_id)

# Initialize SQL execution pipeline for evaluation
pipeline = SQLExecutionPipeline(
    model=config['model']['default'],
    db_type="sqlite"
)

# Run evaluation
evaluation_results = pipeline.evaluate_queries(
    inference_results,
    connection_params={"db_path": config['database']['path']}
)

# Generate and save report
report = pipeline.generate_report(evaluation_results)

# (See run_inference.py in the repository for the complete implementation)

from lamini.experiment.generators import BaseGenerator

class CustomSQLQuestionGenerator(BaseGenerator):
    def __init__(self, model, name):
        instruction = """
        Given the following database schema and glossary, generate {num_variations} 
        variations of each question in the list of questions.

        Schema:
        {schema}

        Glossary:
        {glossary}

        Questions:
        {questions}

        Generate diverse variations that test different SQL concepts while 
        maintaining the same intent as the original questions.
        """

        super().__init__(
            name=name,
            instruction=instruction,
            model=model,
            output_type={"variations": "list"}
        )

    def postprocess(self, prompt):
        # Add any custom post-processing logic
        return prompt

from lamini.experiment.validators import BaseValidator

class CustomSQLValidator(BaseValidator):
    def __init__(self, model, name, db_path):
        instruction = """
        Validate if the following SQL query is correct and will run against 
        the given database schema.

        Schema:
        {schema}

        Query:
        {query}

        Provide detailed feedback on any issues found.
        """

        super().__init__(
            name=name,
            instruction=instruction,
            model=model,
            is_valid_field="is_valid",
            output_type={"is_valid": "bool", "feedback": "str"}
        )

        self.db_path = db_path

    def validate(self, prompt):
        # Add custom validation logic using SQLite
        # Example: Execute query against database and catch errors
        return prompt