Update README.md

README.md

@@ -12,33 +12,25 @@ This is a BERT model fine-tuned on a high-quality Chinese financial dataset. It
 ## 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"
+text = "你是一个聪明的机器人"
+risk_model_name = "risk-model-zh-v0.1"
 
-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)
+risk_tokenizer = AutoTokenizer.from_pretrained(risk_model_name)
+risk_model = AutoModelForSequenceClassification.from_pretrained(risk_model_name)
 
-dataset = load_dataset('json', data_files=dataset_file, cache_dir="cache/", split='train', num_proc=12)
+risk_inputs = risk_tokenizer(text, return_tensors="pt", padding="longest", truncation=True)
+risk_outputs = risk_model(**risk_inputs)
+risk_logits = risk_outputs.logits.squeeze(-1).float().detach().numpy()
 
+risk_score = risk_logits.item()
 
-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()
+result = {
+    "text": text,
+    "risk_score": risk_score
+}
 
-    batch["risk_score"] = logits.tolist()
-    return batch
-
-
-dataset = dataset.map(compute_scores, batched=True, batch_size=512)
-dataset.to_json(output_file)
+print(result)
+# {'text': '你是一个聪明的机器人', 'risk_score': 0.11226219683885574}
 ```