249 lines
7.7 KiB
Python
249 lines
7.7 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Migration script: Convert existing JSON embedding files to FAISS format.
|
|
"""
|
|
|
|
import json
|
|
import pickle
|
|
import re
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
from typing import Any, Dict, List, Optional
|
|
|
|
import faiss
|
|
import numpy as np
|
|
|
|
DATA_DIR = Path("data")
|
|
EMBEDDING_DIM = 1536
|
|
|
|
# Regex patterns from BUGFS_mcp_server.py
|
|
HUNK_HEADER_REGEX = re.compile(
|
|
r"@@ -\d+,\d+ \+(?P<start>\d+),(?P<length>\d+) @@(?:\s+(?P<context>.+))?$",
|
|
flags=re.MULTILINE,
|
|
)
|
|
PYTHON_FUNC_PATTERN = re.compile(r"def\s+(?P<name>\w+)")
|
|
JAVA_METHOD_PATTERN = re.compile(
|
|
r"(?:public|private|protected|static|final|abstract|synchronized|void|\w+)\s+(?P<name>\w+)\s*\("
|
|
)
|
|
CLASS_PATTERN = re.compile(r"(?:class|interface|enum)\s+(?P<name>\w+)")
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class IssueLocation:
|
|
start: int
|
|
length: int
|
|
function: str
|
|
file: str
|
|
|
|
@property
|
|
def label(self) -> str:
|
|
return f"{self.file}:{self.function}"
|
|
|
|
|
|
def normalize_filename(path_str: str) -> str:
|
|
path_str = path_str.replace("\\", "/")
|
|
parts = [p for p in path_str.split("/") if p]
|
|
if not parts:
|
|
return path_str
|
|
return parts[-1]
|
|
|
|
|
|
def extract_method_from_context(context: str) -> Optional[str]:
|
|
if not context:
|
|
return None
|
|
match = PYTHON_FUNC_PATTERN.search(context)
|
|
if match:
|
|
return match.group("name")
|
|
match = JAVA_METHOD_PATTERN.search(context)
|
|
if match:
|
|
name = match.group("name")
|
|
if name not in ("if", "for", "while", "switch", "catch", "synchronized", "return", "new", "throw"):
|
|
return name
|
|
match = CLASS_PATTERN.search(context)
|
|
if match:
|
|
return match.group("name")
|
|
return None
|
|
|
|
|
|
def extract_locations(issue: Dict[str, Any]) -> List[IssueLocation]:
|
|
locations: List[IssueLocation] = []
|
|
seen: set[tuple[str, str]] = set()
|
|
for commit in issue.get("commits", []):
|
|
filename = normalize_filename(commit.get("filename", "unknown"))
|
|
patch = commit.get("patch")
|
|
if not patch:
|
|
continue
|
|
for match in HUNK_HEADER_REGEX.finditer(patch):
|
|
start = int(match.group("start"))
|
|
length = int(match.group("length"))
|
|
context = match.group("context") or ""
|
|
method_name = extract_method_from_context(context)
|
|
if method_name:
|
|
key = (filename, method_name)
|
|
if key not in seen:
|
|
seen.add(key)
|
|
locations.append(
|
|
IssueLocation(start=start, length=length, function=method_name, file=filename)
|
|
)
|
|
return locations
|
|
|
|
|
|
def build_issue_text(issue: Dict[str, Any]) -> str:
|
|
title = (issue.get("title") or "").strip()
|
|
body = (issue.get("body") or "").strip()
|
|
return f"{title}\n\n{body}".strip()
|
|
|
|
|
|
def sanitize_issue_body(issue: Dict[str, Any], min_word_count: int = 1) -> Optional[str]:
|
|
text = build_issue_text(issue)
|
|
text = re.sub(r"\s+", " ", text)
|
|
if len(text.split()) < min_word_count:
|
|
return None
|
|
return text
|
|
|
|
|
|
def migrate_json_to_faiss(json_path: Path) -> None:
|
|
"""Convert a JSON embedding file to FAISS format."""
|
|
print(f"\n{'='*60}")
|
|
print(f"Migrating: {json_path.name}")
|
|
print(f"File size: {json_path.stat().st_size / (1024*1024):.1f} MB")
|
|
print(f"{'='*60}")
|
|
|
|
# Parse filename to get slug and model
|
|
# Format: {slug}_issues_with_{model}.json
|
|
name = json_path.stem # remove .json
|
|
if "_issues_with_" not in name:
|
|
print(f"⚠️ Skipping {json_path.name} - not an embedding file")
|
|
return
|
|
|
|
parts = name.split("_issues_with_")
|
|
slug = parts[0]
|
|
model = parts[1]
|
|
|
|
print(f"Slug: {slug}")
|
|
print(f"Model: {model}")
|
|
|
|
# Load JSON
|
|
print("Loading JSON...")
|
|
with json_path.open("r", encoding="utf-8") as f:
|
|
issues = json.load(f)
|
|
print(f"Loaded {len(issues)} issues")
|
|
|
|
# Build records and embeddings
|
|
print("Processing issues...")
|
|
records: List[Dict[str, Any]] = []
|
|
embeddings: List[np.ndarray] = []
|
|
embeddings_cache: Dict[str, Dict[str, Any]] = {}
|
|
|
|
skipped_no_embedding = 0
|
|
skipped_no_locations = 0
|
|
skipped_no_text = 0
|
|
|
|
for issue in issues:
|
|
embedding = issue.get("embedding")
|
|
if not embedding:
|
|
skipped_no_embedding += 1
|
|
continue
|
|
|
|
# Cache all embeddings (even without locations)
|
|
issue_id = str(issue["id"])
|
|
embeddings_cache[issue_id] = issue
|
|
|
|
locations = extract_locations(issue)
|
|
if not locations:
|
|
skipped_no_locations += 1
|
|
continue
|
|
|
|
text = sanitize_issue_body(issue)
|
|
if not text:
|
|
skipped_no_text += 1
|
|
continue
|
|
|
|
# Convert IssueLocation objects to dicts for pickle compatibility
|
|
locations_as_dicts = [
|
|
{"start": loc.start, "length": loc.length, "function": loc.function, "file": loc.file}
|
|
for loc in locations
|
|
]
|
|
records.append({
|
|
"id": issue["id"],
|
|
"locations": locations_as_dicts,
|
|
"content": text,
|
|
})
|
|
embeddings.append(np.asarray(embedding, dtype=np.float32))
|
|
|
|
print(f" - With embeddings: {len(embeddings_cache)}")
|
|
print(f" - Skipped (no embedding): {skipped_no_embedding}")
|
|
print(f" - Skipped (no locations): {skipped_no_locations}")
|
|
print(f" - Skipped (no text): {skipped_no_text}")
|
|
print(f" - Valid for index: {len(records)}")
|
|
|
|
if not records:
|
|
print("⚠️ No valid records to index!")
|
|
return
|
|
|
|
# Build FAISS index
|
|
print("Building FAISS index...")
|
|
matrix = np.vstack(embeddings).astype(np.float32)
|
|
faiss.normalize_L2(matrix)
|
|
|
|
index = faiss.IndexFlatIP(matrix.shape[1])
|
|
index.add(matrix)
|
|
print(f"Index created with {index.ntotal} vectors, dim={index.d}")
|
|
|
|
# Save files
|
|
faiss_path = DATA_DIR / f"{slug}_{model}.faiss"
|
|
metadata_path = DATA_DIR / f"{slug}_{model}_metadata.pkl"
|
|
embeddings_path = DATA_DIR / f"{slug}_{model}_embeddings.pkl"
|
|
|
|
print(f"Saving FAISS index to {faiss_path.name}...")
|
|
faiss.write_index(index, str(faiss_path))
|
|
|
|
print(f"Saving metadata to {metadata_path.name}...")
|
|
with metadata_path.open("wb") as f:
|
|
pickle.dump(records, f, protocol=pickle.HIGHEST_PROTOCOL)
|
|
|
|
print(f"Saving embeddings cache to {embeddings_path.name}...")
|
|
with embeddings_path.open("wb") as f:
|
|
pickle.dump(embeddings_cache, f, protocol=pickle.HIGHEST_PROTOCOL)
|
|
|
|
# Print size comparison
|
|
old_size = json_path.stat().st_size
|
|
new_size = faiss_path.stat().st_size + metadata_path.stat().st_size + embeddings_path.stat().st_size
|
|
|
|
print(f"\n📊 Size comparison:")
|
|
print(f" - Original JSON: {old_size / (1024*1024):.1f} MB")
|
|
print(f" - FAISS index: {faiss_path.stat().st_size / (1024*1024):.1f} MB")
|
|
print(f" - Metadata: {metadata_path.stat().st_size / (1024*1024):.1f} MB")
|
|
print(f" - Embeddings: {embeddings_path.stat().st_size / (1024*1024):.1f} MB")
|
|
print(f" - Total new: {new_size / (1024*1024):.1f} MB")
|
|
print(f" - Reduction: {(1 - new_size/old_size)*100:.1f}%")
|
|
|
|
print(f"\n✅ Migration complete for {slug}")
|
|
|
|
|
|
def main():
|
|
print("🚀 FAISS Migration Tool")
|
|
print(f"Data directory: {DATA_DIR.absolute()}")
|
|
|
|
# Find all embedding JSON files
|
|
json_files = list(DATA_DIR.glob("*_issues_with_*.json"))
|
|
|
|
if not json_files:
|
|
print("No embedding files found to migrate!")
|
|
return
|
|
|
|
print(f"\nFound {len(json_files)} embedding file(s) to migrate:")
|
|
for f in json_files:
|
|
print(f" - {f.name} ({f.stat().st_size / (1024*1024):.1f} MB)")
|
|
|
|
for json_path in json_files:
|
|
migrate_json_to_faiss(json_path)
|
|
|
|
print("\n" + "="*60)
|
|
print("🎉 All migrations complete!")
|
|
print("="*60)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|