ant-des/pytorch-semantic-dataset-fixed

{
    "input_ids": [2, 847, 3288, ...],                    # Tokenized code (Qwen3-0.6B)
    "semantic_embeddings": [[0.1, -0.2, ...], ...],     # 1024D embeddings per token
    "semantic_positions": [0, 0, 1, 1, 0, ...],         # Binary mask (1=has semantic info)
    "attention_mask": [1, 1, 1, 1, 1, ...],             # Standard attention mask
    "file_path": "torch/nn/modules/linear.py",           # Source file
    "chunk_info": "lines_45_120",                        # Code chunk information
    "num_symbols": 23                                    # Number of semantic symbols
}

# For token "_StreamBase" 
"name: _StreamBase. kind: class_def. type: Type[_StreamBase]. 
 definition: class _StreamBase. description: Base stream class abstraction."

from datasets import load_dataset
from transformers import AutoTokenizer

# Load dataset
dataset = load_dataset("ant-des/pytorch-semantic-dataset-fixed")

# Load corresponding tokenizer
tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-0.6B")

# Access splits
train_dataset = dataset["train"]
eval_dataset = dataset["test"]

print(f"Training samples: {len(train_dataset)}")
print(f"Evaluation samples: {len(eval_dataset)}")

# Get a sample
sample = train_dataset[0]

# Reconstruct the code
code = tokenizer.decode(sample["input_ids"], skip_special_tokens=True)
print("Code:", code[:200] + "...")

# Check semantic coverage
semantic_tokens = sum(sample["semantic_positions"])
total_tokens = len(sample["semantic_positions"])
coverage = semantic_tokens / total_tokens * 100
print(f"Semantic coverage: {coverage:.1f}%")

# Find semantic tokens
for i, (token_id, has_semantic) in enumerate(zip(sample["input_ids"], sample["semantic_positions"])):
    if has_semantic:
        token_text = tokenizer.decode([token_id])
        print(f"Semantic token at position {i}: '{token_text}'")

# Model sees both syntax and semantics
input_ids = [class_token, identifier_token]
semantic_info = [zero_embedding, class_definition_embedding]
# → Better understanding of class structure

class SemanticCodeModel(nn.Module):
    def forward(self, input_ids, semantic_embeddings, semantic_positions):
        # Process both streams
        syntactic_repr = self.language_model(input_ids)
        semantic_repr = self.semantic_projection(semantic_embeddings)
        
        # Cross-attention fusion
        enhanced_repr = self.cross_attention(
            syntactic_repr, semantic_repr, semantic_positions
        )
        return enhanced_repr

# Tree-sitter for precise symbol locations
tree = parser.parse(source_code)
symbols = extract_identifiers(tree)  # Functions, classes, variables

# Jedi for semantic analysis
script = jedi.Script(code=source_code, path=file_path)
definitions = script.goto(line, column, follow_imports=True)
type_info = script.complete(line, column)

def create_semantic_description(symbol_info):
    description = f"name: {symbol.name}. kind: {symbol.type}. type: {symbol.type_hint}."
    
    if symbol.definition:
        description += f" definition: {symbol.definition}."
    
    if symbol.docstring:
        description += f" description: {symbol.docstring[:100]}."
    
    return description

# Generate embeddings
semantic_embeddings = embedding_model.encode(descriptions)

# Align to tokens using Tree-sitter locations
token_embeddings = align_symbols_to_tokens(
    symbols, semantic_embeddings, tokenizer_output
)

from transformers import AutoTokenizer, AutoModel, Trainer, TrainingArguments
import torch.nn as nn

class SemanticCodeModel(nn.Module):
    def __init__(self, base_model_name, semantic_dim=1024):
        super().__init__()
        self.base_model = AutoModel.from_pretrained(base_model_name)
        self.semantic_projection = nn.Linear(semantic_dim, self.base_model.config.hidden_size)
        self.cross_attention = nn.MultiheadAttention(
            self.base_model.config.hidden_size, num_heads=8
        )
        
    def forward(self, input_ids, semantic_embeddings, semantic_positions, **kwargs):
        # Base language model
        outputs = self.base_model(input_ids, **kwargs)
        hidden_states = outputs.last_hidden_state
        
        # Project semantic embeddings
        semantic_proj = self.semantic_projection(semantic_embeddings)
        
        # Apply semantic mask
        masked_semantic = semantic_proj * semantic_positions.unsqueeze(-1)
        
        # Cross-attention fusion
        enhanced_states, _ = self.cross_attention(
            hidden_states, masked_semantic, masked_semantic
        )
        
        return enhanced_states

# Data collator
class SemanticDataCollator:
    def __init__(self, tokenizer):
        self.tokenizer = tokenizer
    
    def __call__(self, batch):
        # Pad sequences and create batch tensors
        max_len = max(len(item["input_ids"]) for item in batch)
        
        # Implement padding logic for all fields
        # ... (see full implementation in repository)

# Training setup
model = SemanticCodeModel("Qwen/Qwen3-0.6B")
tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-0.6B")
data_collator = SemanticDataCollator(tokenizer)

training_args = TrainingArguments(
    output_dir="./semantic-code-model",
    learning_rate=5e-5,
    per_device_train_batch_size=4,
    num_train_epochs=3,
    warmup_steps=100,
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
    data_collator=data_collator,
)

trainer.train()

# Core dependencies
pip install datasets transformers torch
pip install tree-sitter tree-sitter-python
pip install jedi sentence-transformers
pip install numpy huggingface-hub

@dataset{pytorch_semantic_dataset_2024,
    title={PyTorch Semantic Code Dataset: Syntactic Tokenization with Semantic Enrichment},
    author={Antoine Descamps},
    year={2025},
    url={https://huggingface.co/datasets/ant-des/pytorch-semantic-dataset-fixed},
    note={A semantically-enriched Python code dataset combining Tree-sitter and Jedi analysis}
}