Update README.md

README.md

@@ -21,6 +21,8 @@ Gautier Izacard, Mathilde Caron, Lucas Hosseini, Sebastian Riedel, Piotr Bojanow
 
 ## How to use the model
 
+### With the zero-shot classification pipeline
+
 The model can be loaded with the `zero-shot-classification` pipeline like so:
 
 ```python
@@ -44,6 +46,32 @@ candidate_labels = ['travel', 'cooking', 'dancing', 'exploration']
 classifier(sequence_to_classify, candidate_labels, multi_class=True)
 ```
 
+### With manual PyTorch
+
+The model can also be applied on NLI tasks like so:
+
+```python
+import torch
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
+
+# device = "cuda:0" or "cpu"
+device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
+
+model_name = "mjwong/contriever-mnli"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForSequenceClassification.from_pretrained(model_name)
+
+premise = "But I thought you'd sworn off coffee."
+hypothesis = "I thought that you vowed to drink more coffee."
+
+input = tokenizer(premise, hypothesis, truncation=True, return_tensors="pt")
+output = model(input["input_ids"].to(device))
+prediction = torch.softmax(output["logits"][0], -1).tolist()
+label_names = ["entailment", "neutral", "contradiction"]
+prediction = {name: round(float(pred) * 100, 2) for pred, name in zip(prediction, label_names)}
+print(prediction)
+```
+
 ### Eval results
 The model was evaluated using the dev sets for MultiNLI and test sets for ANLI. The metric used is accuracy.