SunDou/dataset-builder

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