Modified model loading

app/main.py

@@ -30,7 +30,7 @@ class BatchResponse(BaseModel):
 # Initialize models
 try:
     logger.info("Loading classification model...")
-    classification_pipeline = CancerClassifier("user1729/BiomedBERT-cancer-bert-classifier-v1.0")
+    classification_pipeline = CancerClassifier()
     
     logger.info("Loading extraction model...")
     extraction_pipeline = CancerExtractor()

app/model.py

@@ -1,12 +1,15 @@
-from transformers import pipeline
+from transformers import AutoModelForSequenceClassification, AutoTokenizer, pipeline
 import os
 import re
 
 class CancerClassifier:
-    def __init__(self, model_path: str):
+    def __init__(self, model_path: "user1729/BiomedBERT-cancer-bert-classifier-v1.0"):
+        model = AutoModelForSequenceClassification.from_pretrained(model_path)
+        tokenizer = AutoTokenizer.from_pretrained(model_path)
         self.classifier = pipeline(
             "text-classification",
-            model=model_path,
+            model=model,
+            tokenizer=tokenizer,
             return_all_scores=True,
             device=0 if os.environ.get("USE_GPU", "false").lower() == "true" else -1,
         )