Added model arch into example

README.md

@@ -33,10 +33,29 @@ The model's bidirectional nature and architectural components make it adept at u
 # Example Code
 
 ```python
+import torch.nn as nn
+import torch.nn.functional as F
+
+class CodeClassifier(nn.Module):
+    def __init__(self, num_classes, embedding_dim, hidden_dim, num_layers, bidirectional=False):
+        super(CodeClassifier, self).__init__()
+        self.feature_extractor = nn.LSTM(embedding_dim, hidden_dim, num_layers, batch_first=True, bidirectional=bidirectional)
+        self.dropout = nn.Dropout(0.5)  # Reintroduce dropout
+        self.fc1 = nn.Linear(hidden_dim * (2 if bidirectional else 1), hidden_dim)  # Intermediate layer
+        self.fc2 = nn.Linear(hidden_dim, num_classes)  # Output layer
+
+    def forward(self, x):
+        x = x.unsqueeze(1)  # Add sequence dimension
+        x, _ = self.feature_extractor(x)
+        x = x.squeeze(1)  # Remove sequence dimension
+        x = self.fc1(x)
+        x = self.dropout(x)  # Apply dropout
+        x = self.fc2(x)
+        return x
+
 import torch
 from transformers import AutoTokenizer, AutoModel
 from pathlib import Path
-from model import CodeClassifier
 
 def infer(text, model_path, embedding_model_name):
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")