Update app.py

app.py

@@ -25,8 +25,8 @@ tokenizer = AutoTokenizer.from_pretrained(model_name,use_fast=False)
 model = AutoModelForSequenceClassification.from_pretrained(model_name)
 
 
-premise = sentence1
-hypothesis = sentence2
+premise = selected_sentence1
+hypothesis = selected_sentence2
 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()
@@ -35,14 +35,10 @@ prediction = {name: round(float(pred) * 100, 1) for pred, name in zip(prediction
 print(prediction)
 
 
-from sentence_transformers import CrossEncoder
-model1 = CrossEncoder('cross-encoder/nli-deberta-v3-xsmall')
-scores1 = model1.predict([(sentence1, sentence2)])
-
-#Convert scores to labels
-label_mapping = ['contradiction', 'entailment', 'neutral']
-labels = [label_mapping[score_max] for score_max in scores1.argmax(axis=1)]
-labels
+from transformers import pipeline
+pipe = pipeline("text-classification",model="sileod/deberta-v3-base-tasksource-nli")
+labels=pipe([dict(text=selected_sentence1,
+  text_pair=selected_sentence2)])
 
 
 
@@ -67,8 +63,8 @@ def extract_person_names(sentence):
     
     return person_names[0]
 
-person_name1 = extract_person_names(sentence1)
-person_name2 = extract_person_names(sentence2)
+person_name1 = extract_person_names(selected_sentence1)
+person_name2 = extract_person_names(selected_sentence2)
 
 st.write("Result:", prediction)