Update README.md

README.md

@@ -11,6 +11,34 @@ pipeline_tag: text-classification
 
 ## Model Details
 
+```python
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+device = "cuda" # the device to load the model onto
+
+model = AutoModelForCausalLM.from_pretrained("armaniii/llama-argument-classification")
+tokenizer = AutoTokenizer.from_pretrained("armaniii/lllama-argument-classification")
+
+model.to(device)
+model.eval()
+
+for batch in tqdm.tqdm(data):
+    with torch.no_grad():
+        input_text = tokenizer(batch, padding=True, truncation=True,max_length=2048,return_tensors="pt").to(device)
+        output = model(**input_text)
+        logits = output.logits
+        predicted_class = torch.argmax(logits, dim=1)
+        # Convert logits to a list of predicted labels
+        predictions.extend(predicted_class.cpu().tolist())
+
+    # Get the ground truth labels
+df["predictions"] = predictions
+
+num2label = {
+0:"NoArgument",
+1:"Argument"
+}
+```
 ### Model Description
 
 <!-- Provide a longer summary of what this model is. -->