Update README.md

README.md

@@ -22,18 +22,14 @@ from transformers import AutoTokenizer, AutoModelForSequenceClassification
 import torch
 
 device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
-model_name = "nightdessert
-/
-WeCheck "
+model_name = "nightdessert/WeCheck"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModelForSequenceClassification.from_pretrained(model_name)
-premise = "I first thought that I liked the movie, but upon second thought it was actually disappointing."
-hypothesis = "The movie was not good."
+premise = "I first thought that I liked the movie, but upon second thought it was actually disappointing." # Input for Summarization/ Dialogue / Paraphrase
+hypothesis = "The movie was not good."   # Output for Summarization/ Dialogue / Paraphrase
 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()
-label_names = ["entailment", "neutral", "contradiction"]
-prediction = {name: round(float(pred) * 100, 1) for pred, name in zip(prediction, label_names)}
+output = model(input["input_ids"].to(device))[:,0]  # device = "cuda:0" or "cpu"
+prediction = torch.sigmoid(output).tolist()
 print(prediction)
 
 license: openrail