Update app.py to load model from Hugging Face Model Hub

app.py

@@ -10,24 +10,19 @@ class CodeClassifier(torch.nn.Module):
     def __init__(self, base_model, num_labels=6):
         super(CodeClassifier, self).__init__()
         self.base = base_model
-        self.reduction = torch.nn.Linear(768, 512)  # Randomly initialized
-        self.classifier = torch.nn.Linear(512, num_labels)  # Match checkpoint
+        self.reduction = torch.nn.Linear(768, 512)
+        self.classifier = torch.nn.Linear(512, num_labels)
 
     def forward(self, input_ids, attention_mask):
         outputs = self.base(input_ids=input_ids, attention_mask=attention_mask)
         reduced = self.reduction(outputs.pooler_output)
         return self.classifier(reduced)
 
-# Load model and tokenizer
+# Load model and tokenizer from Hugging Face Model Hub
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 tokenizer = RobertaTokenizer.from_pretrained('microsoft/codebert-base')
-base_model = RobertaModel.from_pretrained('microsoft/codebert-base')
+base_model = RobertaModel.from_pretrained('martynattakit/CodeSentinel-Model')  # Match your Model repo
 model = CodeClassifier(base_model)
-checkpoint = torch.load("C:\\Users\\MartyNattakit\\Downloads\\best_model.pt", map_location=device)
-
-# Load the state dict, focusing on classifier weights
-model_state = checkpoint.get('model_state_dict', checkpoint)
-model.load_state_dict(model_state, strict=False)
 print("Loaded state dict keys:", model.state_dict().keys())
 print("Classifier weight shape:", model.classifier.weight.shape)
 model.eval()