updated readme: corrected the model usage code

README.md

@@ -43,25 +43,15 @@ import torch
 import torch.nn.functional as F
 from transformers import AutoModelForSequenceClassification, AutoTokenizer
 
-# Load ModernBERT NLI model and tokenizer
-MODEL_ID = "dleemiller/ModernCE-base-nli"
+MODEL_ID = "ChandruK/MarkOne"
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
 model = AutoModelForSequenceClassification.from_pretrained(MODEL_ID).to(device)
 model.eval()
 
-# Dynamically resolve Entailment label index from config
-label_mapping = model.config.id2label
-ENTAILMENT_IDX = None
-for idx, label_name in label_mapping.items():
-    if label_name.lower() in ["entailment", "entail"]:
-        ENTAILMENT_IDX = idx
-        break
-if ENTAILMENT_IDX is None:
-    ENTAILMENT_IDX = 1
-
-# 1. Define Candidate Labels and Hypothesis Template
+ENTAILMENT_IDX = 1
+
 CANDIDATE_LABELS = [
     "a Word document creation request",
     "an Excel spreadsheet creation request",
@@ -72,14 +62,11 @@ CANDIDATE_LABELS = [
 ]
 HYPOTHESIS_TEMPLATE = "This user prompt is {}."
 
-# 2. Input Prompt
 prompt = "Generate a formal project proposal in Word format."
 
-# 3. Construct Premise-Hypothesis Pairs
 hypotheses = [HYPOTHESIS_TEMPLATE.format(label) for label in CANDIDATE_LABELS]
 premises = [prompt] * len(CANDIDATE_LABELS)
 
-# 4. Tokenize Pairs
 inputs = tokenizer(
     text=premises,
     text_pair=hypotheses,
@@ -89,20 +76,17 @@ inputs = tokenizer(
     return_tensors="pt",
 ).to(device)
 
-# 5. Model Inference
 with torch.no_grad():
     outputs = model(**inputs)
     logits = outputs.logits
     entailment_logits = logits[:, ENTAILMENT_IDX] if logits.shape[-1] >= 3 else logits[:, 1]
 
-    # Softmax probability normalization with temperature scaling
     probabilities = F.softmax(entailment_logits / 0.5, dim=-1)
     top_prob, top_idx = torch.max(probabilities, dim=-1)
 
 predicted_label = CANDIDATE_LABELS[top_idx.item()]
 confidence = top_prob.item() * 100
 
-# Build dictionary containing key details and all category probabilities
 label_scores = {
     label: round(prob.item() * 100, 2)
     for label, prob in zip(CANDIDATE_LABELS, probabilities)
@@ -115,6 +99,5 @@ output_json_data = {
     "all_scores": label_scores,
 }
 
-# Output as formatted JSON string
 json_output = json.dumps(output_json_data, indent=2)
 print(json_output)