tricoder 1.2.8__py3-none-any.whl

tricoder/context_view.py

@@ -0,0 +1,228 @@
+"""Context view: Node2Vec-style random walks and Word2Vec."""
+import random
+from multiprocessing import Pool, cpu_count
+from typing import List, Tuple
+
+import numpy as np
+from gensim.models import Word2Vec
+from gensim.models.keyedvectors import KeyedVectors
+
+
+def _get_num_workers() -> int:
+    """Get number of workers (all cores - 1, minimum 1)."""
+    return max(1, cpu_count() - 1)
+
+
+def _generate_walks_for_node(args):
+    """Generate walks for a single node (helper for multiprocessing)."""
+    start_node, adj_list, num_walks, walk_length, p, q, seed_offset = args
+    # Set seed for this worker
+    random.seed(42 + seed_offset + start_node)
+    np.random.seed(42 + seed_offset + start_node)
+
+    walks = []
+    for walk_idx in range(num_walks):
+        if not adj_list[start_node]:
+            continue
+
+        walk = [start_node]
+
+        for _ in range(walk_length - 1):
+            curr = walk[-1]
+            neighbors = adj_list[curr]
+
+            if not neighbors:
+                break
+
+            if len(walk) == 1:
+                next_node = random.choice(neighbors)[0]
+            else:
+                prev = walk[-2]
+                probs = []
+                nodes = []
+
+                for neighbor, weight in neighbors:
+                    nodes.append(neighbor)
+                    if neighbor == prev:
+                        prob = weight / p
+                    elif neighbor in adj_list[prev]:
+                        prob = weight
+                    else:
+                        prob = weight / q
+                    probs.append(max(prob, 1e-10))
+
+                probs = np.array(probs)
+                probs = probs / probs.sum()
+                next_node = np.random.choice(nodes, p=probs)
+
+            walk.append(next_node)
+
+        walks.append([str(node) for node in walk])
+
+    return walks
+
+
+def generate_random_walks(edges: List[Tuple[int, int, str, float]],
+                          num_nodes: int,
+                          num_walks: int = 10,
+                          walk_length: int = 80,
+                          p: float = 1.0,
+                          q: float = 1.0,
+                          random_state: int = 42,
+                          n_jobs: int = -1) -> List[List[int]]:
+    """
+    Generate Node2Vec-style random walks with multiprocessing support.
+    
+    Args:
+        edges: list of (src_idx, dst_idx, relation, weight) tuples
+        num_nodes: number of nodes
+        num_walks: number of walks per node
+        walk_length: length of each walk
+        p: return parameter (1/p controls likelihood of returning)
+        q: in-out parameter (1/q controls likelihood of going further)
+        random_state: random seed
+        n_jobs: number of parallel jobs (-1 for all cores - 1)
+    
+    Returns:
+        List of walks (each walk is a list of node indices)
+    """
+    # Build adjacency list
+    adj_list = {i: [] for i in range(num_nodes)}
+    for src_idx, dst_idx, rel, weight in edges:
+        adj_list[src_idx].append((dst_idx, weight))
+        adj_list[dst_idx].append((src_idx, weight))
+
+    # Determine number of workers
+    if n_jobs == -1:
+        n_jobs = _get_num_workers()
+    n_jobs = max(1, min(n_jobs, num_nodes))
+
+    # Prepare arguments for parallel processing
+    nodes_to_process = [node for node in range(num_nodes) if adj_list[node]]
+
+    if n_jobs == 1 or len(nodes_to_process) < n_jobs:
+        # Sequential processing for small cases
+        random.seed(random_state)
+        np.random.seed(random_state)
+        walks = []
+        for _ in range(num_walks):
+            for start_node in nodes_to_process:
+                walk = [start_node]
+                for _ in range(walk_length - 1):
+                    curr = walk[-1]
+                    neighbors = adj_list[curr]
+                    if not neighbors:
+                        break
+
+                    if len(walk) == 1:
+                        next_node = random.choice(neighbors)[0]
+                    else:
+                        prev = walk[-2]
+                        probs = []
+                        nodes = []
+                        for neighbor, weight in neighbors:
+                            nodes.append(neighbor)
+                            if neighbor == prev:
+                                prob = weight / p
+                            elif neighbor in adj_list[prev]:
+                                prob = weight
+                            else:
+                                prob = weight / q
+                            probs.append(max(prob, 1e-10))
+                        probs = np.array(probs)
+                        probs = probs / probs.sum()
+                        next_node = np.random.choice(nodes, p=probs)
+                    walk.append(next_node)
+                walks.append([str(node) for node in walk])
+    else:
+        # Parallel processing
+        args_list = [
+            (node, adj_list, num_walks, walk_length, p, q, random_state + i)
+            for i, node in enumerate(nodes_to_process)
+        ]
+
+        # Parallel processing - use standard Pool which works on all platforms
+        chunksize = max(1, len(args_list) // (n_jobs * 2))
+        with Pool(processes=n_jobs) as pool:
+            results = pool.map(_generate_walks_for_node, args_list, chunksize=chunksize)
+
+        # Flatten results
+        walks = []
+        for result in results:
+            walks.extend(result)
+
+    return walks
+
+
+def train_word2vec(walks: List[List[str]], dim: int, window: int = 7,
+                   negative: int = 3, epochs: int = 3, random_state: int = 42,
+                   n_jobs: int = -1, batch_words: int = 10000) -> KeyedVectors:
+    """
+    Train Word2Vec SkipGram model on random walks with multiprocessing.
+    Optimized defaults: window=7 (was 10), negative=3 (was 5) for faster training.
+    
+    Args:
+        walks: list of walks (each walk is a list of node ID strings)
+        dim: embedding dimensionality
+        window: context window size (reduced default: 7)
+        negative: number of negative samples (reduced default: 3)
+        epochs: number of training epochs
+        random_state: random seed
+        n_jobs: number of parallel workers (-1 for all cores - 1)
+        batch_words: words per batch (larger = faster but more memory)
+    
+    Returns:
+        Trained KeyedVectors model
+    """
+    if n_jobs == -1:
+        n_jobs = _get_num_workers()
+
+    # For gensim 4.x, use workers parameter
+    workers = max(1, n_jobs)
+    
+    model = Word2Vec(
+        sentences=walks,
+        vector_size=dim,
+        window=window,
+        min_count=1,
+        workers=workers,
+        sg=1,  # SkipGram
+        negative=negative,
+        epochs=epochs,
+        seed=random_state,
+        batch_words=batch_words  # Larger batches for faster training
+    )
+
+    return model.wv
+
+
+def compute_context_view(edges: List[Tuple[int, int, str, float]],
+                         num_nodes: int,
+                         dim: int,
+                         num_walks: int = 10,
+                         walk_length: int = 80,
+                         random_state: int = 42,
+                         n_jobs: int = -1) -> Tuple[np.ndarray, KeyedVectors]:
+    """
+    Compute context view embeddings using Node2Vec + Word2Vec with multiprocessing.
+    
+    Returns:
+        embeddings: node embeddings from context view
+        keyed_vectors: Word2Vec KeyedVectors model
+    """
+    walks = generate_random_walks(edges, num_nodes, num_walks, walk_length,
+                                  random_state=random_state, n_jobs=n_jobs)
+    kv = train_word2vec(walks, dim, random_state=random_state, n_jobs=n_jobs)
+
+    # Extract embeddings for all nodes
+    embeddings = np.zeros((num_nodes, dim))
+    for i in range(num_nodes):
+        node_str = str(i)
+        if node_str in kv:
+            embeddings[i] = kv[node_str]
+        else:
+            # Initialize with small random values if node not seen
+            np.random.seed(random_state + i)
+            embeddings[i] = np.random.normal(0, 0.01, dim)
+
+    return embeddings, kv

tricoder/data_loader.py

@@ -0,0 +1,144 @@
+"""Data loading utilities for TriVector Code Intelligence."""
+import json
+from collections import defaultdict
+from typing import Dict, List, Tuple
+
+from .subtoken_utils import extract_subtokens, get_file_hierarchy
+
+
+def load_nodes(nodes_path: str) -> Tuple[
+    Dict[str, int], List[Dict], Dict[str, List[str]], Dict[str, Tuple[str, str, str]]]:
+    """
+    Load nodes from JSONL file.
+    
+    Returns:
+        node_to_idx: mapping from node ID to index
+        node_metadata: list of node metadata dictionaries
+        node_subtokens: mapping from node_id to list of normalized subtokens
+        node_file_info: mapping from node_id to (file_name, directory_path, top_level_package)
+    """
+    node_to_idx = {}
+    node_metadata = []
+    node_subtokens = {}
+    node_file_info = {}
+
+    with open(nodes_path, 'r') as f:
+        for line in f:
+            if not line.strip():
+                continue
+            node = json.loads(line)
+            node_id = node['id']
+            if node_id not in node_to_idx:
+                idx = len(node_to_idx)
+                node_to_idx[node_id] = idx
+
+                # Extract name and metadata
+                name = node.get('name', '')
+                meta = node.get('meta', {})
+
+                # Extract subtokens
+                raw_subtokens, normalized_subtokens = extract_subtokens(name, normalize=True)
+                node_subtokens[node_id] = normalized_subtokens
+
+                # Extract file hierarchy
+                file_path = meta.get('file', '') if isinstance(meta, dict) else ''
+                file_name, directory_path, top_level_package = get_file_hierarchy(file_path)
+                node_file_info[node_id] = (file_name, directory_path, top_level_package)
+
+                # Store metadata with subtoken info for debugging
+                meta_with_subtokens = meta.copy() if isinstance(meta, dict) else {}
+                meta_with_subtokens['_raw_subtokens'] = raw_subtokens
+                meta_with_subtokens['_normalized_subtokens'] = normalized_subtokens
+
+                node_metadata.append({
+                    'id': node_id,
+                    'kind': node.get('kind', 'unknown'),
+                    'name': name,
+                    'meta': meta_with_subtokens
+                })
+
+    return node_to_idx, node_metadata, node_subtokens, node_file_info
+
+
+def load_edges(edges_path: str, node_to_idx: Dict[str, int]) -> Tuple[List[Tuple[int, int, str, float]], int]:
+    """
+    Load edges from JSONL file.
+    
+    Returns:
+        edges: list of (src_idx, dst_idx, relation, weight) tuples
+        num_nodes: number of unique nodes
+    """
+    edges = []
+    seen_nodes = set()
+
+    with open(edges_path, 'r') as f:
+        for line in f:
+            if not line.strip():
+                continue
+            edge = json.loads(line)
+            src_id = edge['src']
+            dst_id = edge['dst']
+
+            if src_id in node_to_idx and dst_id in node_to_idx:
+                src_idx = node_to_idx[src_id]
+                dst_idx = node_to_idx[dst_id]
+                rel = edge.get('rel', 'unknown')
+                weight = float(edge.get('weight', 1.0))
+                edges.append((src_idx, dst_idx, rel, weight))
+                seen_nodes.add(src_idx)
+                seen_nodes.add(dst_idx)
+
+    num_nodes = len(node_to_idx)
+    return edges, num_nodes
+
+
+def load_types(types_path: str, node_to_idx: Dict[str, int]) -> Tuple[
+    Dict[int, Dict[str, int]], Dict[str, int]]:
+    """
+    Load type tokens from JSONL file.
+    
+    Returns:
+        node_types: mapping from node_idx to {type_token: count}
+        type_to_idx: mapping from type token to index
+    """
+    node_types = defaultdict(lambda: defaultdict(int))
+    type_to_idx = {}
+
+    with open(types_path, 'r') as f:
+        for line in f:
+            if not line.strip():
+                continue
+            entry = json.loads(line)
+            symbol_id = entry['symbol']
+            type_token = entry['type_token']
+            count = int(entry.get('count', 1))
+
+            if symbol_id in node_to_idx:
+                node_idx = node_to_idx[symbol_id]
+                node_types[node_idx][type_token] += count
+
+                if type_token not in type_to_idx:
+                    type_to_idx[type_token] = len(type_to_idx)
+
+    return dict(node_types), type_to_idx
+
+
+def build_node_location_map(node_metadata: List[Dict]) -> Dict[int, Tuple[str, int]]:
+    """
+    Build mapping from node_idx to (file_path, line_number) for context window co-occurrence.
+    
+    Args:
+        node_metadata: list of node metadata dictionaries
+    
+    Returns:
+        Mapping from node_idx to (file_path, line_number)
+    """
+    location_map = {}
+    for idx, node_meta in enumerate(node_metadata):
+        meta = node_meta.get('meta', {})
+        if isinstance(meta, dict):
+            file_path = meta.get('file', '')
+            lineno = meta.get('lineno', -1)
+            if file_path and lineno >= 0:
+                location_map[idx] = (file_path, lineno)
+    return location_map