| # 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 |
|
|
