update to new model with weighted loss 0.1 for class 0

config.json

@@ -57,7 +57,6 @@
   "num_hidden_layers": 12,
   "pad_token_id": 0,
   "position_embedding_type": "absolute",
-  "problem_type": "single_label_classification",
   "torch_dtype": "float32",
   "transformers_version": "4.41.2",
   "type_vocab_size": 2,

model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:39cb42cf548de528ab2d1505a7ba84031e250ba2d092ddd861ad109ad0acbcb3
+oid sha256:7556885b937a337a064f088a19141cd29ef6e6f2276cf53b70b1a1730b2c99d4
 size 439832632

sdg_predict/cli_predict.py

@@ -6,31 +6,53 @@ from tqdm import tqdm
 import sys
 import torch
 from sdg_predict.inference import load_model, predict
+import logging
+
+# Set up logging
+logging.basicConfig(
+    level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s", force=True
+)
+
 
 def main():
-    parser = argparse.ArgumentParser(description="Batch inference using Hugging Face model.")
+    parser = argparse.ArgumentParser(
+        description="Batch inference using Hugging Face model."
+    )
     parser.add_argument("input", type=Path, help="Input JSONL file")
-    parser.add_argument("--key", type=str, required=True, help="JSON key with text input")
-    parser.add_argument("--batch_size", type=int, default=8, help="Batch size")
-    parser.add_argument("--model", type=str, default="simon-clmtd/sdg-scibert-zo_up", help="Model name on the Hub")
-    parser.add_argument("--top1", action="store_true", help="Return only top prediction")
-    parser.add_argument("--output", type=Path, help="Output file (optional, otherwise stdout)")
+    parser.add_argument(
+        "--key", type=str, default="text", help="JSON key with text input"
+    )
+    parser.add_argument("--batch_size", "-b", type=int, default=8, help="Batch size")
+    parser.add_argument(
+        "--model",
+        type=str,
+        default="simon-clmtd/sdg-scibert-zo_up",
+        help="Model name on the Hub",
+    )
+    parser.add_argument(
+        "--top1", action="store_true", help="Return only top prediction"
+    )
+    parser.add_argument(
+        "--output", type=Path, help="Output file (optional, otherwise stdout)"
+    )
     args = parser.parse_args()
 
     # -------------------------------
-# 1. Device Setup (MPS support for Apple Silicon)
-# -------------------------------
+    # 1. Device Setup (MPS support for Apple Silicon)
+    # -------------------------------
     if torch.backends.mps.is_available():
         device = torch.device("mps")
-        print("Using MPS device")
+        logging.info("Using MPS device")
     elif torch.cuda.is_available():
         device = torch.device("cuda")
-        print("Using CUDA device")
+        logging.info("Using CUDA device")
     else:
         device = torch.device("cpu")
-        print("Using CPU device")
-        
+        logging.info("Using CPU device")
+    # device = torch.device("cpu")
+    logging.info("Loading model: %s", args.model)
     tokenizer, model = load_model(args.model, device)
+    logging.info("Model loaded successfully")
 
     with args.input.open() as f:
         texts = []
@@ -40,20 +62,44 @@ def main():
             if args.key not in row:
                 continue
             texts.append(row[args.key])
+            logging.debug("Text: %s", row[args.key])
             rows.append(row)
 
+    logging.info("Starting predictions on %d texts", len(texts))
     predictions = predict(
         texts,
         tokenizer,
         model,
         device,
         batch_size=args.batch_size,
-        return_all_scores=not args.top1
+        return_all_scores=not args.top1,
     )
+    logging.info("Predictions completed")
 
     output_stream = args.output.open("w") if args.output else sys.stdout
     for row, pred in zip(rows, predictions):
-        row["prediction"] = pred
-        print(json.dumps(row, ensure_ascii=False), file=output_stream)
+        # Compute binary probabilities for labels 1-17
+        binary_predictions = {}
+        for label_data in pred:
+            label_data["score"] = round(
+                label_data["score"], 3
+            )  # Round prediction scores to 3 decimal places
+            label = int(label_data["label"])
+            if 1 <= label <= 17:
+                binary_prob = label_data["score"]  # Already rounded
+                binary_predictions[str(label)] = binary_prob
+
+        output_row = {
+            "id": row.get("id"),
+            "text": row.get("text"),
+            "prediction": pred,
+            "binary_predictions": binary_predictions,
+        }
+        print(json.dumps(output_row, ensure_ascii=False), file=output_stream)
     if args.output:
         output_stream.close()
+        logging.info("Output written to %s", args.output)
+
+
+if __name__ == "__main__":
+    main()

sdg_predict/inference.py

@@ -1,42 +1,45 @@
 # sdg_predict/inference.py
-from transformers import AutoTokenizer, AutoModelForSequenceClassification
+from transformers import AutoTokenizer, AutoModelForSequenceClassification, pipeline
 import torch
+import logging
+
 
 def load_model(model_name, device):
-    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    tokenizer = AutoTokenizer.from_pretrained(model_name, do_lower_case=False)
     model = AutoModelForSequenceClassification.from_pretrained(model_name).to(device)
     model.eval()
     return tokenizer, model
 
+
 def batched(iterable, batch_size):
     for i in range(0, len(iterable), batch_size):
-        yield iterable[i:i + batch_size]
+        yield iterable[i : i + batch_size]
+
 
 def predict(texts, tokenizer, model, device, batch_size=8, return_all_scores=True):
-    results = []
-    for batch_texts in batched(texts, batch_size):
-        inputs = tokenizer(
-            batch_texts,
-            return_tensors="pt",
-            padding=True,
-            truncation=True,
-            max_length=512
-        ).to(device)
-
-        with torch.no_grad():
-            logits = model(**inputs).logits
-            probs = torch.nn.functional.softmax(logits, dim=-1)
-
-        for prob in probs:
-            if return_all_scores:
-                results.append([
-                    {"label": model.config.id2label[i], "score": prob[i].item()}
-                    for i in range(len(prob))
-                ])
-            else:
-                top = torch.argmax(prob).item()
-                results.append({
-                    "label": model.config.id2label[top],
-                    "score": prob[top].item()
-                })
+    classifier = pipeline(
+        "text-classification",
+        model=model,
+        tokenizer=tokenizer,
+        device=device,
+        batch_size=batch_size,
+        truncation=True,
+        padding=True,
+        max_length=512,
+        top_k=None if return_all_scores else 1,
+    )
+
+    results = classifier(texts)
+    if return_all_scores:
+        for result in results:
+            for score in result:
+                score["score"] = round(
+                    score["score"], 3
+                )  # Round scores to 3 decimal places
+    else:
+        for result in results:
+            result["score"] = round(
+                result["score"], 3
+            )  # Round top score to 3 decimal places
+
     return results

setup.py

@@ -2,7 +2,7 @@ from setuptools import setup, find_packages
 
 setup(
     name="sdg-predict",
-    version="0.1",
+    version="0.2",
     packages=find_packages(),
     install_requires=[
         "transformers>=4.36",

training_args.bin

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:a6728d45aae9302ac2656ddf218d44224848150f29ea3fafc5535dde9cd98003
+oid sha256:a4b473744ec4c80646022813576aa0fa32733d67c31a15bc75b51c2d5cb456e6
 size 5713