Update README.md

README.md

@@ -45,33 +45,19 @@ This model is a fine-tuned version of `roberta-large` trained for logical fallac
 To use the model for quick classification with a text pipeline:
 
 ```python
-from transformers import pipeline, AutoTokenizer, AutoModelForSequenceClassification
+from transformers import pipeline
 import torch
 
-# Replace with your Hugging Face model path
+# Initialize the text classification pipeline with the specified model and tokenizer
 model_path = "MidhunKanadan/roberta-large-fallacy-classification"
+pipe = pipeline("text-classification", model=model_path, tokenizer=model_path, use_fast=False, device=0 if torch.cuda.is_available() else -1)
 
-# Automatically detect the device (GPU if available, else CPU)
-device = 0 if torch.cuda.is_available() else -1
-
-# Initialize the text classification pipeline with use_fast=False for the tokenizer
-tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=False)
-model = AutoModelForSequenceClassification.from_pretrained(model_path)
-
-# Use GPU if available
-pipe = pipeline("text-classification", model=model, tokenizer=tokenizer, device=device)
-
-# Define a sample text
+# Sample text to analyze for logical fallacies
 text = "The rooster crows always before the sun rises, therefore the crowing rooster causes the sun to rise."
+result = pipe(text)[0]  # Retrieve the first result (main prediction)
 
-# Make a prediction
-result = pipe(text)
-
-# Print the predicted label and score
-print(f"Predicted Label: {result[0]['label']}")
-print(f"Score: {result[0]['score']:.4f}")
-
-
+# Output the predicted label and confidence score
+print(f"Predicted Label: {result['label']}\nScore: {result['score']:.4f}")
 ```
 
 Expected Output:
@@ -88,33 +74,20 @@ For more control, load the model and tokenizer directly and perform classificati
 import torch
 from transformers import AutoTokenizer, AutoModelForSequenceClassification
 
-# Load your model and tokenizer from Hugging Face
+# Load model and tokenizer, set device
 model_path = "MidhunKanadan/roberta-large-fallacy-classification"
-model = AutoModelForSequenceClassification.from_pretrained(model_path)
-tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=False)
-
-# Automatically move model to GPU if available, else use CPU
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-model.to(device)
+model = AutoModelForSequenceClassification.from_pretrained(model_path).to(device).eval()
+tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=False)
 
-# Define a sample text
+# Tokenize input and get probabilities
 text = "The rooster crows always before the sun rises, therefore the crowing rooster causes the sun to rise."
-
-# Tokenize the input text and move tensors to the appropriate device
 inputs = tokenizer(text, return_tensors="pt").to(device)
-
-# Run the model and get logits
 with torch.no_grad():
-    logits = model(**inputs).logits
-
-# Apply softmax to get probabilities
-probabilities = torch.nn.functional.softmax(logits, dim=1)[0]
-
-# Pair each label with its score and sort in descending order
-sorted_results = sorted(zip(model.config.id2label.values(), probabilities), key=lambda x: x[1], reverse=True)
+    probabilities = torch.nn.functional.softmax(model(**inputs).logits, dim=1)[0]
 
-# Print each label and its corresponding score in descending order
-for label, score in sorted_results:
+# Output sorted results
+for label, score in sorted(zip(model.config.id2label.values(), probabilities), key=lambda x: x[1], reverse=True):
     print(f"{label}: {score.item():.4f}")
 ```