Spaces:
Sleeping
Sleeping
File size: 3,083 Bytes
bca5172 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 | """
Simple prediction demo: load a fine-tuned HF model and classify cleaned files.
Usage (bash):
python scripts/predict_demo.py --model_dir "./darkbert_finetuned_final" --input data/cleaned/cleaned_extracted_text.jsonl --out predictions.jsonl
"""
import argparse
import json
import os
import torch
import numpy as np
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch.nn.functional as F
def predict_text_single(model, tokenizer, text, device='cpu', max_length=512, use_prompt=False, format_prompt_fn=None):
inp = format_prompt_fn(text) if use_prompt and format_prompt_fn is not None else text
enc = tokenizer(inp, truncation=True, padding=True, return_tensors='pt', max_length=max_length).to(device)
with torch.no_grad():
outputs = model(**enc)
probs = F.softmax(outputs.logits, dim=-1).cpu().numpy()[0]
pred = int(np.argmax(probs))
return pred, float(probs[pred]), probs.tolist()
def read_jsonl(path):
with open(path, 'r', encoding='utf-8') as f:
for line in f:
yield json.loads(line)
def main():
parser = argparse.ArgumentParser()
parser.add_argument('--model_dir', required=True)
parser.add_argument('--input', required=True)
parser.add_argument('--out', default='predictions.jsonl')
parser.add_argument('--device', default='cpu')
parser.add_argument('--use_prompt', action='store_true')
args = parser.parse_args()
predict_file(args.model_dir, args.input, args.out, device=args.device, use_prompt=args.use_prompt)
def predict_file(model_dir: str, input_path: str, out_path: str, device: str = 'cpu', use_prompt: bool = False):
"""Programmatic API: load model from `model_dir`, predict on `input_path` JSONL, write predictions to `out_path`.
Returns the number of predictions written.
"""
tokenizer = AutoTokenizer.from_pretrained(model_dir)
model = AutoModelForSequenceClassification.from_pretrained(model_dir).to(device)
# try to import format_prompt function from notebook code if available
try:
from finetune_prompt_template import format_prompt_for_training as fmt
except Exception:
fmt = None
count = 0
os.makedirs(os.path.dirname(out_path) or '.', exist_ok=True)
with open(out_path, 'w', encoding='utf-8') as out_f:
for obj in read_jsonl(input_path):
text = obj.get('text', '')
pred, conf, probs = predict_text_single(model, tokenizer, text, device=device, use_prompt=use_prompt, format_prompt_fn=fmt)
out = {
'text': text,
'pred_id': pred,
'confidence': conf,
'probs': probs
}
if 'label' in obj:
out['label'] = obj['label']
out_f.write(json.dumps(out, ensure_ascii=False) + '\n')
count += 1
print(f"Predictions saved to {out_path} ({count} items)")
return count
if __name__ == '__main__':
main()
|