File size: 1,640 Bytes
a528044 210d151 a528044 |
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 |
---
license: apache-2.0
language:
- zh
pipeline_tag: text-classification
library_name: transformers
---
# risk-model-zh-v0.1
## Introduction
This is a BERT model fine-tuned on a high-quality Chinese financial dataset. It generates a security risk score, which helps to identify and remove data with security risks from financial datasets, thereby reducing the proportion of illegal or undesirable data. For the complete data cleaning process, please refer to [YiZhao](https://github.com/HITsz-TMG/YiZhao).
## Quickstart
Here is an example code snippet for generating security risk scores using this model.
```python
import torch
from datasets import load_dataset
from transformers import AutoTokenizer, AutoModelForSequenceClassification
model_name = "risk-model-zh-v0.1"
dataset_file = "your_dataset.jsonl"
text_column = "text"
output_file = "your_output.jsonl"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSequenceClassification.from_pretrained(model_name, torch_dtype=torch.bfloat16)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)
dataset = load_dataset('json', data_files=dataset_file, cache_dir="cache/", split='train', num_proc=12)
def compute_scores(batch):
inputs = tokenizer(batch[text_column], return_tensors="pt", padding="longest", truncation=True).to(device)
with torch.no_grad():
outputs = model(**inputs)
logits = outputs.logits.squeeze(-1).float().cpu().numpy()
batch["risk_score"] = logits.tolist()
return batch
dataset = dataset.map(compute_scores, batched=True, batch_size=512)
dataset.to_json(output_file)
``` |