the-jashthakkar/CodeMode

1'''2Synthesize triplet and positive pair datasets from chunked code files.'''3 4import argparse5import json6import random7import hashlib8from pathlib import Path9from typing import Dict, List10from datetime import datetime11from sklearn.feature_extraction.text import TfidfVectorizer12from sklearn.metrics.pairwise import cosine_similarity13 14 15# ============================16# CONFIG17# ============================18 19MAX_DOCUMENTS = 20020POSITIVE_VARIANTS = 521TFIDF_MAX_FEATURES = 500022RANDOM_SEED = 4223 24BASE_OUTPUT_DIR = Path("data/synthetic")25 26random.seed(RANDOM_SEED)27 28 29# ============================30# UTILITIES31# ============================32 33def load_chunks(file_path):34    path = Path(file_path)35 36    if path.suffix == ".jsonl":37        chunks = []38        with open(path, "r", encoding="utf-8") as f:39            for line_no, line in enumerate(f, 1):40                line = line.strip()41                if not line:42                    continue43                try:44                    chunks.append(json.loads(line))45                except json.JSONDecodeError as e:46                    raise ValueError(47                        f"Invalid JSON on line {line_no} in {path}"48                    ) from e49        return chunks50 51    elif path.suffix == ".json":52        with open(path, "r", encoding="utf-8") as f:53            data = json.load(f)54        if not isinstance(data, list):55            raise ValueError(f"{path} must contain a list of chunks")56        return data57 58    else:59        raise ValueError(60            f"Unsupported file format {path.suffix}. Use .json or .jsonl"61        )62 63 64 65def save_jsonl(path: Path, records: List[Dict]):66    path.parent.mkdir(parents=True, exist_ok=True)67    with path.open("w", encoding="utf-8") as f:68        for r in records:69            f.write(json.dumps(r, ensure_ascii=False) + "\n")70 71 72def save_json(path: Path, data):73    path.parent.mkdir(parents=True, exist_ok=True)74    with path.open("w", encoding="utf-8") as f:75        json.dump(data, f, indent=2)76 77 78def stable_document_id(chunk: Dict, idx: int) -> str:79    """80    Generate a canonical, stable document_id.81    """82    base = f"{chunk.get('file_path','unknown')}::{idx}"83    return "doc_" + hashlib.sha1(base.encode()).hexdigest()84 85 86def infer_framework(input_path: Path) -> str:87    """88    Infer framework from path (fallback-safe).89    """90    parts = [p.lower() for p in input_path.parts]91    for fw in ["crewai", "langchain", "langgraph", "autogen"]:92        if fw in parts:93            return fw94    return "unknown"95 96 97# ============================98# ANCHOR GENERATION (LLM PLACEHOLDER)99# ============================100 101def generate_anchor_questions(code: str, n: int) -> List[str]:102    """103    Deterministic placeholder (LLM-ready).104    """105    symbol = code.split("(")[0].replace("def ", "").replace("class ", "").strip()106 107    templates = [108        f"How does {symbol} work in Python?",109        f"How to implement {symbol}?",110        f"Example usage of {symbol}",111        f"Explain the {symbol} logic",112        f"Best practices for {symbol}",113    ]114 115    random.shuffle(templates)116    return templates[:n]117 118 119# ============================120# NEGATIVE MINING121# ============================122 123def build_tfidf(chunks: List[Dict]):124    corpus = [c["code"] for c in chunks]125    vectorizer = TfidfVectorizer(126        stop_words="english",127        max_features=TFIDF_MAX_FEATURES128    )129    matrix = vectorizer.fit_transform(corpus)130    return vectorizer, matrix131 132 133def mine_hard_negative(134    anchor: str,135    positive_idx: int,136    chunks: List[Dict],137    vectorizer,138    matrix,139) -> Dict:140    query_vec = vectorizer.transform([anchor])141    scores = cosine_similarity(query_vec, matrix)[0]142 143    ranked = sorted(144        [(i, s) for i, s in enumerate(scores)],145        key=lambda x: x[1],146        reverse=True,147    )148 149    for idx, _ in ranked:150        if idx != positive_idx:151            return chunks[idx]152 153    raise RuntimeError("No negative candidate found")154 155 156# ============================157# MAIN PIPELINE158# ============================159 160def generate_datasets(input_path: Path, run_name: str):161    output_dir = BASE_OUTPUT_DIR / run_name162    framework = infer_framework(input_path)163 164    chunks = load_chunks(input_path)165    # Filter only semantic code chunks166    chunks = [167        c for c in chunks168        if c.get("chunk_type") in {"class", "method", "function"}169        and "code" in c170    ]171 172    random.shuffle(chunks)173    chunks = chunks[:MAX_DOCUMENTS]174 175    # Assign canonical document_id176    for idx, c in enumerate(chunks):177        c["document_id"] = stable_document_id(c, idx)178 179    vectorizer, matrix = build_tfidf(chunks)180 181    positive_pairs = []182    triplets = []183 184    for idx, chunk in enumerate(chunks):185        code = chunk["code"]186        doc_id = chunk["document_id"]187 188        # -------- POSITIVE PAIRS --------189        anchors = generate_anchor_questions(code, POSITIVE_VARIANTS)190        for a in anchors:191            positive_pairs.append({192                "document_id": doc_id,193                "anchor": a,194                "positive": code,195                "framework": framework,196                "source": "synthetic_positive_v2",197            })198 199        # -------- TRIPLET --------200        anchor = anchors[0]201        negative_chunk = mine_hard_negative(202            anchor, idx, chunks, vectorizer, matrix203        )204 205        triplets.append({206            "document_id": doc_id,207            "anchor": anchor,208            "positive": code,209            "negative": negative_chunk["code"],210            "framework": framework,211            "source": "synthetic_triplet_v2",212        })213 214    # -------- SAVE --------215    save_jsonl(output_dir / "positive_pairs.jsonl", positive_pairs)216    save_jsonl(output_dir / "triplets.jsonl", triplets)217 218    save_json(output_dir / "positive_pairs.json", positive_pairs)219    save_json(output_dir / "triplets.json", triplets)220 221    metadata = {222        "name": run_name,223        "framework": framework,224        "input_file": str(input_path),225        "num_chunks": len(chunks),226        "positive_pairs": len(positive_pairs),227        "triplets": len(triplets),228        "created_at": datetime.utcnow().isoformat(),229        "random_seed": RANDOM_SEED,230    }231 232    save_json(output_dir / "metadata.json", metadata)233 234    print(f"✅ Dataset generated at: {output_dir}")235 236 237# ============================238# ENTRY POINT239# ============================240 241if __name__ == "__main__":242    parser = argparse.ArgumentParser()243    parser.add_argument("--input", required=True, help="Chunked JSONL file")244    parser.add_argument("--name", required=True, help="Synthetic dataset name")245 246    args = parser.parse_args()247 248    generate_datasets(249        input_path=Path(args.input),250        run_name=args.name,251    )252 253# # For document id 254 255# document_id := sha1(256#     normalized_repo_path +257#     file_path +258#     top_level_symbol259# )