Add CLI and inference modules for batch prediction using Hugging Face model

sdg_predict/cli_predict.py

@@ -0,0 +1,47 @@
+# sdg_predict/cli_predict.py
+import argparse
+import json
+from pathlib import Path
+from tqdm import tqdm
+import sys
+import torch
+from sdg_predict.inference import load_model, predict
+
+def main():
+    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)")
+    args = parser.parse_args()
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    tokenizer, model = load_model(args.model, device)
+
+    with args.input.open() as f:
+        texts = []
+        rows = []
+        for line in f:
+            row = json.loads(line)
+            if args.key not in row:
+                continue
+            texts.append(row[args.key])
+            rows.append(row)
+
+    predictions = predict(
+        texts,
+        tokenizer,
+        model,
+        device,
+        batch_size=args.batch_size,
+        return_all_scores=not args.top1
+    )
+
+    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)
+    if args.output:
+        output_stream.close()

sdg_predict/inference.py

@@ -0,0 +1,42 @@
+# sdg_predict/inference.py
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
+import torch
+
+def load_model(model_name, device):
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    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]
+
+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()
+                })
+    return results

setup.py

@@ -0,0 +1,20 @@
+from setuptools import setup, find_packages
+
+setup(
+    name="sdg-predict",
+    version="0.1",
+    packages=find_packages(),
+    install_requires=[
+        "transformers>=4.36",
+        "torch>=2.0",
+        "tqdm",
+    ],
+    entry_points={
+        "console_scripts": [
+            "sdg-predict = sdg_predict.cli_predict:main"
+        ]
+    },
+    author="Simon Clematide",
+    description="Command-line prediction for SDG SciBERT classifier",
+    python_requires=">=3.8",
+)