Novora/CodeClassifier-v1-Tiny

1---2license: apache-2.03datasets:4    - Novora/CodeClassifier_v1 5pipeline_tag: text-classification 6---7 8# Introduction9 10Novora Code Classifier v1 Tiny, is a tiny `Text Classification` model, which classifies given code text input under 1 of `31` different classes (programming languages).11 12This model is designed to be able to run on CPU, but optimally runs on GPUs.13 14# Info15- 1 of 31 classes output16- 512 token input dimension17- 64 hidden dimensions18- 2 linear layers19- The `snowflake-arctic-embed-xs` model is used as the embeddings model.20- Dataset split into 80% training set, 20% testing set.21- The combined test and training data is around 1000 chunks per programming language, the data is 31,100 chunks (entries) as 512 tokens per chunk, being a snippet of the code.22- Picked from the 18th epoch out of 20 done.23 24# Architecture25 26The `CodeClassifier-v1-Tiny` model employs a neural network architecture optimized for text classification tasks, specifically for classifying programming languages from code snippets. This model includes:27 28- **Bidirectional LSTM Feature Extractor**: This bidirectional LSTM layer processes input embeddings, effectively capturing contextual relationships in both forward and reverse directions within the code snippets.29 30- **Fully Connected Layers**: The network includes two linear layers. The first projects the pooled features into a hidden feature space, and the second linear layer maps these to the output classes, which correspond to different programming languages. A dropout layer with a rate of 0.5 between these layers helps mitigate overfitting.31 32The model's bidirectional nature and architectural components make it adept at understanding the syntax and structure crucial for code classification.33 34# Testing/Training Datasets35I have put here the samples entered into the training/testing pipeline, its a very small amount.36 37| Language     | Testing Count | Training Count |38|--------------|---------------|----------------|39| Ada          | 20            | 80             |40| Assembly     | 20            | 80             |41| C            | 20            | 80             |42| C#           | 20            | 80             |43| C++          | 20            | 80             |44| COBOL        | 14            | 55             |45| Common Lisp  | 20            | 80             |46| Dart         | 20            | 80             |47| Erlang       | 20            | 80             |48| F#           | 20            | 80             |49| Go           | 20            | 80             |50| Haskell      | 20            | 80             |51| Java         | 20            | 80             |52| JavaScript   | 20            | 80             |53| Julia        | 20            | 80             |54| Kotlin       | 20            | 80             |55| Lua          | 20            | 80             |56| MATLAB       | 20            | 80             |57| PHP          | 20            | 80             |58| Perl         | 20            | 80             |59| Prolog       | 1             | 4              |60| Python       | 20            | 80             |61| R            | 20            | 80             |62| Ruby         | 20            | 80             |63| Rust         | 20            | 80             |64| SQL          | 20            | 80             |65| Scala        | 20            | 80             |66| Swift        | 20            | 80             |67| TypeScript   | 20            | 80             |68 69# Example Code70 71```python72import torch.nn as nn73import torch.nn.functional as F74 75class CodeClassifier(nn.Module):76    def __init__(self, num_classes, embedding_dim, hidden_dim, num_layers, bidirectional=False):77        super(CodeClassifier, self).__init__()78        self.feature_extractor = nn.LSTM(embedding_dim, hidden_dim, num_layers, batch_first=True, bidirectional=bidirectional)79        self.dropout = nn.Dropout(0.5)  # Reintroduce dropout80        self.fc1 = nn.Linear(hidden_dim * (2 if bidirectional else 1), hidden_dim)  # Intermediate layer81        self.fc2 = nn.Linear(hidden_dim, num_classes)  # Output layer82 83    def forward(self, x):84        x = x.unsqueeze(1)  # Add sequence dimension85        x, _ = self.feature_extractor(x)86        x = x.squeeze(1)  # Remove sequence dimension87        x = self.fc1(x)88        x = self.dropout(x)  # Apply dropout89        x = self.fc2(x)90        return x91 92import torch93from transformers import AutoTokenizer, AutoModel94from pathlib import Path95 96def infer(text, model_path, embedding_model_name):97    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")98    99    # Load tokenizer and embedding model100    tokenizer = AutoTokenizer.from_pretrained(embedding_model_name)101    embedding_model = AutoModel.from_pretrained(embedding_model_name).to(device)102    embedding_model.eval()103 104    # Prepare inputs105    inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True)106    inputs = {k: v.to(device) for k, v in inputs.items()}107    108    # Generate embeddings109    with torch.no_grad():110        embeddings = embedding_model(**inputs)[0][:, 0]111 112    # Load classifier model113    model = CodeClassifier(num_classes=31, embedding_dim=embeddings.size(-1), hidden_dim=64, num_layers=2, bidirectional=True)114    model.load_state_dict(torch.load(model_path, map_location=device))115    model = model.to(device)116    model.eval()117 118    # Predict class119    with torch.no_grad():120        output = model(embeddings)121        _, predicted = torch.max(output, dim=1)122 123    # Language labels124    languages = [125        'Ada', 'Assembly', 'C', 'C#', 'C++', 'COBOL', 'Common Lisp', 'Dart', 'Erlang', 'F#',126        'Fortran', 'Go', 'Haskell', 'Java', 'JavaScript', 'Julia', 'Kotlin', 'Lua', 'MATLAB',127        'Objective-C', 'PHP', 'Perl', 'Prolog', 'Python', 'R', 'Ruby', 'Rust', 'SQL', 'Scala',128        'Swift', 'TypeScript'129    ]130    131    return languages[predicted.item()]132 133# Example usage134if __name__ == "__main__":135    example_text = "print('Hello, world!')"  # Replace with actual text for inference136    model_file_path = Path("./model.safetensors")137    predicted_language = infer(example_text, model_file_path, "Snowflake/snowflake-arctic-embed-xs")138    print(f"Predicted programming language: {predicted_language}")139 140```141