Update README.md

README.md

@@ -15,34 +15,25 @@ To collect training samples, we use the **Qwen-72B** model to thoroughly annotat
 ## Quickstart
 Here is an example code snippet for generating financial relevance scores using this model.
 ```python
-import torch
-from datasets import load_dataset
 from transformers import AutoTokenizer, AutoModelForSequenceClassification
 
-model_name = "fin-model-en-v0.1"
-dataset_file = "your_dataset.jsonl"
-text_column = "text"
-output_file = "your_output.jsonl"
+text = "You are a smart robot"
+fin_model_name = "fin-model-en-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)
+fin_tokenizer = AutoTokenizer.from_pretrained(fin_model_name)
+fin_model = AutoModelForSequenceClassification.from_pretrained(fin_model_name)
 
-dataset = load_dataset('json', data_files=dataset_file, cache_dir="cache/", split='train', num_proc=12)
+fin_inputs = fin_tokenizer(text, return_tensors="pt", padding="longest", truncation=True)
+fin_outputs = fin_model(**fin_inputs)
+fin_logits = fin_outputs.logits.squeeze(-1).float().detach().numpy()
 
+fin_score = fin_logits.item()
+result = {
+    "text": text,
+    "fin_score": fin_score,
+    "fin_int_score": int(round(max(0, min(fin_score, 5))))
+}
 
-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["fin_score"] = logits.tolist()
-    batch["fin_int_score"] = [int(round(max(0, min(score, 5)))) for score in logits]
-    return batch
-
-
-dataset = dataset.map(compute_scores, batched=True, batch_size=512)
-dataset.to_json(output_file)
+print(result)
+# {'text': 'You are a smart robot', 'fin_score': 0.3258197605609894, 'fin_int_score': 0}
 ```