Update README.md

README.md

@@ -34,22 +34,13 @@ You'll need to pip install transformers & maybe sentencepiece
 
 ### How to use
 ```python
-from transformers import AutoTokenizer, AutoModelForSequenceClassification
-import torch, time
-device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
+from transformers import pipeline
+device = 0 if torch.cuda.is_available() else -1
 model_name = 'Reggie/muppet-roberta-base-joke_detector'
 max_seq_len = 510
 
-tokenizer = AutoTokenizer.from_pretrained(model_name, model_max_length=max_seq_len)
-model = AutoModelForSequenceClassification.from_pretrained(model_name).to(device)
-
-premise = """A nervous passenger is about to book a flight ticket, and he asks the airlines' ticket seller, "I hope your planes are safe. Do they have a good track record for safety?" The airline agent replies, "Sir, I can guarantee you, we've never had a plane that has crashed more than once." """
-hypothesis = ""
-
-input = tokenizer(premise, hypothesis, truncation=True, return_tensors="pt")
-output = model(input["input_ids"].to(device))  # device = "cuda:0" or "cpu"
-prediction = torch.softmax(output["logits"][0], -1).tolist()
-is_joke = True if prediction[0] < prediction[1] else False
-
-print(is_joke)
+pipe = pipeline(model=model_name, device=device, truncation=True, max_length=max_seq_len)
+is_it_a_joke = """A nervous passenger is about to book a flight ticket, and he asks the airlines' ticket seller, "I hope your planes are safe. Do they have a good track record for safety?" The airline agent replies, "Sir, I can guarantee you, we've never had a plane that has crashed more than once." """
+result = pipe(is_it_a_joke) # [{'label': 'LABEL_1', 'score': 0.7313136458396912}]
+print('This is a joke') if result[0]['label'] == 'LABEL_1' else print('This is not a joke')
 ```