Update app.py

app.py

@@ -8,6 +8,7 @@ tokenizer_name = "allenai/scibert_scivocab_uncased"
 tokenizer = AutoTokenizer.from_pretrained(tokenizer_name)
 model = AutoModelForSequenceClassification.from_pretrained(model_name)
 
+
 def inference(abstract: str):
     """
     Split an abstract into sentences and perform claim identification.
@@ -20,14 +21,13 @@ def inference(abstract: str):
         truncation=True,
         padding="longest"
     )
-    output = model(**inputs).logits
-    for (idx, out) in enumerate(output):
-        pred = out.argmax().item()
-        if pred:
-            claims.append(sents[idx])
+    logits = model(**inputs).logits
+    preds = logits.argmax(dim=1)  # convert logits to predictions
+    claims = [sent for sent, pred in zip(sents, preds) if pred == 1]
     if len(claims) > 0:
         return ".\n".join(claims)
-    return "No claims were made here"
+    else:
+        return "No claims found from a given abstract."
 
 
 with gr.Blocks(theme=gr.themes.Soft()) as demo: