megatran

KO | EN

gitlite — search

#python #ai #llm #go #golang #machine-learning #deep-learning #linux #docker #awesome #kubernetes #vue #security #open-source

(★ 11)

[VLDB'25] Official repo for Paper "Weak-to-Strong Prompts with Lightweight-to-Powerful LLMs for High-Accuracy, Low-Cost, and Explainable Data Transformation"

megatran Latest Version Download

Download Latest Version (.zip)

// repository documentation

Was this content helpful?

★ 0 (0 ratings)

# MegaTran

[![Python](https://img.shields.io/badge/Python-3.10+-3776AB)](https://www.python.org/downloads/release/python-3100/)
[![Read Paper](https://img.shields.io/badge/Read_Paper-PVLDB-B22222)](https://www.vldb.org/pvldb/vol18/p2371-tang.pdf) 
[![HuggingFace Model](https://img.shields.io/badge/HuggingFace-Model-FF8800)](https://huggingface.co/Ti-ger/llama3_lora_dt_chat)
[![Citation](https://img.shields.io/badge/Cite_Us-BibTex-25D366)](https://raw.githubusercontent.com/HKUSTDial/megatran/main/)

## Framework Overview
![Framework](https://raw.githubusercontent.com/HKUSTDial/megatran/main/img/works.png)

MegaTran is a **two-stage** framework that generates robust data transformation code from simple user requests. It bridges the gap between user intent and high-quality code by first strengthening the prompt and then using an optimized code generation process.

**Stage 1: Weak2StrongPrompt**
- A lightweight LLM interprets the user's intent to create a strong prompt.
- How it works: 
    - We use a fine-tuned, lightweight LLM that is specifically trained to analyze a user's request. It predicts the correct transformation "operator" (e.g., format, extract) and generates a detailed natural language description of the task.
    - The fine-tuned model is avaliable for download at HuggingFace.

**Stage 2: Prompt2Code**
- A powerful LLM generates code based on the strong prompt, guided by two key optimizations for error correction and knowledge retrieval.
- How it works:
    - **Sanity-Check Reflection**: An iterative self-correction mechanism that system checks code against a pre-defined while extensible checklist of common errors (e.g. syntax, value).
    - **Lazy-RAG**: When an error occurs with a unknown library, it queries a pre-built vector database that contains documentation and code snippets from sites such as GitHub and PyPI.

## Setup

1. Clone this repo and install dependencies
```bash
pip install -r requirements.txt
```

2. Configure environment variables
```bash
# Create .env file with your API keys
OPENAI_API_KEY=your_api_key_here
```

3. Start `vLLM` server for hosting the fine-tuned model
```bash
vllm serve \
    --model ./assets/models/llama3_lora_sft \ # wait for HF downloading ...
    --config ./etc/vllm-server.yaml
```
> Note: You can use `CUDA_VISIBLE_DEVICES` to target the GPU device for vLLM server

4. Test Weak2Strong module inference
```bash
python w2s_prompt_inference.py -q "input:abc, output:ABC"

# Expected output: 
# format(): Convert the string to uppercase
```

5. [offline, optional] Build RAG vector database
```bash
# Build vector database for code libs retrieval
python scripts/build_vector_db.py \
    --config etc/vec_db.yaml \
    [-q "hijri date to gregorian date"] # test single query by adding this argument
```
> A pre-built vector database is saved in `assets/rag/code_db`

## User Guide

1. Execute the main pipeline
```bash
# Test mode (with smaller dataset)
python run.py \
    --config etc/mega-transform.yaml \
    --exp_name demo \
    --model gpt-4o-mini \
    --testing

# Full dataset run
python run.py \
    --config etc/mega-transform.yaml \
    --exp_name exp-1 \
    --model gpt-4o-mini \
    --dataset_name stackoverflow
```

2. Check experiment results as show in `demo` folder. Results will include:
- Code Generation Results (per task)
- Full test results (full_result.csv)
- Summary statistics (task-level accuracy, token usage, etc.)
- Runtime logs for current run

## Project Structure
```
chat-transform/
├── run.py                # Main execution script
├── w2s_prompt_inference.py # Weak2strong prompt inference
├── etc/                  # Configuration files
│   ├── mega-transform.yaml # pipeline config
│   ├── code-llm.yaml       # baseline Code LLM
│   ├── vllm-server.yaml    # vLLM server config
│   └── vec_db.yaml         # RAG vector database config
├── framework/            # Core components
│   ├── chat_to_inst.py   # Chat to instruction conversion
│   ├── code_generator.py # Code generation
│   ├── lazy_rag.py       # Lazy RAG module
│   ├── reflection.py     # Sanity-check Reflection module
│   └── prompt_generator.py # Prompt composition
├── util/                 # Utility modules
│   ├── analyzer.py       # Result analysis and reporting
│   ├── load_data.py      # Data loading utilities
│   ├── context_manager.py # Context management
│   └── __init__.py
├── assets/               # Model assets
│   ├── models/           # Fine-tuned models
│   └── rag/              # RAG related files (Vec DB, list of missing packages)
├── scripts/              # Utility scripts
│   ├── build_vector_db.py # Build RAG vector database
│   ├── foundation_model.py # Foundation model baseline
│   └── push_to_hf.py      # Push to HF
├── temp/                 # Temporary files (on-the-fly generated code)
├── .env                  # Environment variables
└── requirements.txt      # Project dependencies
```

## Baselines
Foundation model baseline, source code refer to  the orginal implementation [here](https://github.com/HazyResearch/fm_data_tasks/blob/main/notebooks/data_transformation_experiments.ipynb)
```bash
# Dataset: benchmark-stackoverflow
python scripts/foundation_model.py --dataset stackoverflow --model gpt-4o-mini

# Dataset: benchmark-BinqQuery (semantic)
python scripts/foundation_model.py --dataset bingquery-logs --model gpt-4o-mini
```

Naive code generation baseline:
```bash
python run.py \
    --config etc/code-llm.yaml \ # use code-llm config here
    --exp_name exp-1 \
    --model gpt-4o-mini \
    --dataset_name stackoverflow
```

## Citation

If you find this work useful, please cite:

```bibtex
@article{DBLP:journals/pvldb/LiYLFT25,
  author       = {Changlun Li and
                  Chenyu Yang and
                  Yuyu Luo and
                  Ju Fan and
                  Nan Tang},
  title        = {Weak-to-Strong Prompts with Lightweight-to-Powerful LLMs for High-Accuracy,
                  Low-Cost, and Explainable Data Transformation},
  journal      = {Proc. {VLDB} Endow.},
  volume       = {18},
  number       = {8},
  pages        = {2371--2384},
  year         = {2025}
}
```

megatran

megatran Latest Version Download

Today's picks

Was this content helpful?

Recent Feedback

megatran

megatran Latest Version Download

Today's picks

Was this content helpful?

Recent Feedback

Download README