Update README.md

README.md

@@ -12,7 +12,7 @@ Model was trained with a `max_length` of `4096`, but the base model supports `81
 ### Example Code
 ```py
 import torch
-from transformers import AutoTokenizer, LlamaForSequenceClassification
+from transformers import pipeline
 import json
 from tqdm import tqdm
 
@@ -45,34 +45,19 @@ def load_json_or_jsonl(file_path):
         return None
 
 
-tokenizer = AutoTokenizer.from_pretrained(
-    "PJMixers/Danube3-ClassTest-v0.1-500M"
-)
-model = LlamaForSequenceClassification.from_pretrained(
-    "PJMixers/Danube3-ClassTest-v0.1-500M",
+pipe = pipeline(
+    task="text-classification",
+    model="PJMixers/Danube3-ClassTest-v0.1-500M",
     device_map="cuda",
-    torch_dtype=torch.bfloat16,
-    attn_implementation="sdpa",
+    torch_dtype=torch.bfloat16
 )
-
 data = load_json_or_jsonl(
     "./PrefMix-Classifier-Data-validation.json"
 )
 
 passes, fails = 0, 0
 for sample in tqdm(data):
-    input_text = sample["input_text"]
-    true_label = sample["labels"]
-
-    inputs = tokenizer(
-        input_text,
-        return_tensors="pt"
-    ).to("cuda")
-
-    with torch.no_grad():
-        generated_label = model(**inputs).logits.argmax().item()
-
-    if generated_label == true_label:
+    if int(pipe(sample["input_text"])[0]["label"]) == sample["labels"]:
         passes += 1
     else:
         fails += 1