Spaces:
Sleeping
Sleeping
| """ | |
| 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() | |