# Dataset Builder

This project contains the complete source code for building three domain-specific datasets for scientific computing code intelligence research.

## Project Structure

```
dataset_builder/
├── README.md                          # This file
│
├── data1/                             # DATA1: Domain-Specific Code Dataset
│   ├── main.py                        # Step 0-1: Keyword expansion + GitHub repo search
│   ├── main_v2.py                     # Step 0-4: Full pipeline (search → check → clone → filter)
│   ├── util.py                        # Shared utilities (logger, LLM calls, code extensions)
│   ├── download_dataset.py            # Download ChemPile code dataset from HuggingFace
│   ├── merge_dataset.py               # Merge crawled repos with ChemPile data, deduplicate
│   ├── analysis.py                    # Code-level analysis (comments, functions, tokens)
│   ├── compute_stars_keywords.py      # Compute stars/keyword statistics
│   ├── compute_statistics.py          # Compute code statistics from JSONL analysis files
│   ├── rename.py                      # Rename repo directories to owner___repo format
│   ├── rename2.py                     # Rename ChemPile files with zero-padded numbering
│   ├── pyproject.toml                 # Python project config
│   ├── scripts/
│   │   └── export_files_to_csv.py     # Export repo files to CSV grouped by keyword
│   ├── reporting/                     # Statistical reporting and visualization
│   │   ├── __init__.py
│   │   ├── main.py                    # Reporting entry point
│   │   ├── visualization.py           # Generate figures (funnel, distributions, etc.)
│   │   ├── repo_meta_scan.py          # Scan repo-level metadata
│   │   ├── code_file_stats.py         # File-level code statistics
│   │   ├── code_file_stats_fast.py    # Optimized file-level statistics
│   │   ├── stage_a_stats.py           # Stage A (search/check) statistics
│   │   ├── stage_b_stats.py           # Stage B (clone/filter) statistics
│   │   └── join_insights.py           # Join and cross-analyze insights
│   └── README.md                      # DATA1 dataset documentation
│
├── data2/                             # DATA2: Code-Documentation Alignment Dataset
│   ├── instruction_generation/        # README summarization pipeline
│   │   ├── pipeline.py                # Unified entry (summarize + parse modes)
│   │   ├── summarize_repo_readme.py   # Summarize repo READMEs using LLM
│   │   ├── extract_repo_functions.py  # Extract functions from repos
│   │   ├── schemas.py                 # Pydantic data schemas
│   │   └── prompts/
│   │       ├── function_extract.txt   # Prompt for function extraction
│   │       └── readme_summary.txt     # Prompt for README summarization
│   ├── step22/                        # Function scoring, generation, alignment
│   │   ├── build.py                   # Build tree-sitter language parsers
│   │   ├── func_stat.py               # Extract functions using tree-sitter
│   │   ├── md_stat.py                 # Extract & save README summaries
│   │   ├── emb_qwen_func.py           # Score functions using Qwen embedding model
│   │   ├── emb_qwen_md.py             # Score READMEs using Qwen embedding model
│   │   ├── function_req.py            # Filter functions by score threshold
│   │   ├── gemini_generation.py       # Generate docstrings using Gemini API
│   │   ├── alignment.py               # Align functions with generated docstrings
│   │   ├── prompt.txt                 # Prompt template for docstring generation
│   │   ├── depend_analysis.py         # Dependency/call-graph analysis
│   │   ├── find_none_score_func.py    # Find functions missing scores
│   │   ├── folder_stat.py             # Repository folder statistics
│   │   ├── ppt.py                     # Visualization of alignment data
│   │   └── debug_parser.py            # Debug tree-sitter parser loading
│   └── README.md                      # DATA2 dataset documentation
│
├── data3/                             # DATA3: Programming Problems Generation Dataset
│   ├── main.py                        # RepoAgent: generate docs for repos
│   ├── gemini.py                      # Gemini API connectivity test
│   ├── load_dataset.py                # Load and inspect datasets
│   ├── instruct_generation.py         # Score functions for scientific relevance
│   ├── extract_functions.py           # Extract functions from enhanced_dataset.csv
│   ├── extract_functions_v2.py        # Extract functions v2 (better CSV/JSON handling)
│   ├── merge_datasets.py              # Merge res2.csv with dataset_all.csv
│   ├── generate_programming_problems.py  # Generate problems using Gemini API
│   ├── generate_problems_batch.py     # Batch problem generation (OpenAI batch API)
│   ├── generate_problems_openai.py    # Problem generation via OpenAI API
│   ├── enrich_programming_problems.py # Enrich problems with source code context
│   ├── vllm_high.py                   # VLLM-based high-throughput inference
│   ├── vllm_qwen_batch.py            # Qwen model batch inference via VLLM
│   ├── show_pricing.py                # Display API pricing information
│   ├── check_enhanced.py              # Validate enhanced dataset
│   ├── check_index_distribution.py    # Check index distribution
│   ├── check_match.py                 # Check data matching
│   ├── check_relationship.py          # Check data relationships
│   ├── is_sci_prompt.txt              # Prompt: classify code as scientific computing
│   ├── is_sci_prompt1.txt             # Prompt variant for scientific classification
│   ├── score_prompt.txt               # Prompt: score function relevance
│   ├── *.sh                           # Various shell scripts for batch processing
│   └── README.md                      # DATA3 dataset documentation
```

## Dataset Building Pipelines

### DATA1: Domain-Specific Code Dataset

**Goal**: Collect, filter, and export domain-specific code from GitHub repositories.

**Pipeline** (executed in order):

1. **Keyword Expansion & Search** (`main.py` / `main_v2.py`)
   - Expand scientific keywords using LLM
   - Search GitHub API for repositories matching keywords
   - Check relevance using LLM (reads READMEs)
   - Clone relevant repos (shallow clone)
   - Filter to keep only code files

2. **External Data** (`download_dataset.py`)
   - Download ChemPile code dataset from HuggingFace

3. **Merge & Deduplicate** (`merge_dataset.py`)
   - Merge crawled repos with ChemPile data
   - Deduplicate by content hash

4. **Analysis** (`analysis.py`, `compute_*.py`)
   - Analyze code metrics (lines, comments, functions, tokens)
   - Compute keyword and stars statistics

5. **Export** (`scripts/export_files_to_csv.py`)
   - Export final dataset to CSV files grouped by keyword

6. **Reporting** (`reporting/`)
   - Generate statistical reports and visualizations

### DATA2: Code-Documentation Alignment Dataset

**Goal**: Generate high-quality docstrings for scientific code functions.

**Pipeline** (executed in order):

1. **README Summarization** (`instruction_generation/`)
   - Summarize repository READMEs using LLM
   - Extract structured information from repos

2. **Function Extraction** (`step22/func_stat.py`)
   - Parse code using tree-sitter to extract functions
   - Multi-language support (Python, C, C++, Java, Go, Rust, Julia)

3. **README Processing** (`step22/md_stat.py`)
   - Copy README summaries to function dataset directories

4. **Embedding Scoring** (`step22/emb_qwen_func.py`, `emb_qwen_md.py`)
   - Score function quality using Qwen embedding model
   - Score README quality using Qwen embedding model

5. **Function Filtering** (`step22/function_req.py`)
   - Filter functions by combined quality score

6. **Docstring Generation** (`step22/gemini_generation.py`)
   - Generate docstrings using Gemini API
   - Budget monitoring with circuit breaker
   - Checkpoint/resume support

7. **Alignment** (`step22/alignment.py`)
   - Merge function data with generated docstrings

### DATA3: Programming Problems Generation Dataset

**Goal**: Generate programming problems inspired by scientific code.

**Pipeline** (executed in order):

1. **Documentation Generation** (`main.py`)
   - Use RepoAgent to generate documentation for repositories

2. **Function Extraction** (`extract_functions.py`, `extract_functions_v2.py`)
   - Extract individual functions from enhanced dataset

3. **Scientific Relevance Scoring** (`instruct_generation.py`)
   - Score functions for scientific computing relevance
   - Use `is_sci_prompt.txt` and `score_prompt.txt` as prompts

4. **Dataset Merge** (`merge_datasets.py`)
   - Merge function scores with source code data

5. **Problem Generation** (`generate_programming_problems.py`)
   - Generate programming problems using Gemini API
   - Filter by relevance score
   - Budget monitoring and cost control

6. **Enrichment** (`enrich_programming_problems.py`)
   - Enrich generated problems with source code context

## Dependencies

### Common
- `pandas`, `tqdm`, `jsonlines`
- `python-dotenv`

### DATA1
- `langchain`, `langchain-openai`, `pydantic`, `loguru`
- `requests` (GitHub API)
- `matplotlib`, `seaborn`, `wordcloud` (reporting)
- `datasets` (HuggingFace)

### DATA2
- `tree-sitter`, `tree-sitter-python`, `tree-sitter-c`, etc.
- `vllm`, `transformers`, `torch` (embedding scoring)
- `google-cloud-aiplatform`, `vertexai` (Gemini API)

### DATA3
- `google-cloud-aiplatform`, `vertexai` (Gemini API)
- `openai` (OpenAI API)
- `vllm`, `transformers`, `torch` (local inference)

## Notes

- Scripts contain hardcoded paths that need to be updated for your environment
- API credentials (GitHub token, Gemini, OpenAI) need to be configured separately
- Large datasets require significant storage and compute resources
- Most scripts support checkpoint/resume for long-running processes