Fix model combination

app.py

@@ -1,3 +1,4 @@
+import torch
 import streamlit as st
 from datasets import load_dataset
 from transformers import DistilBertForSequenceClassification, DistilBertTokenizer
@@ -34,15 +35,20 @@ def predict(model_abstract, model_claims, tokenizer_abstract, tokenizer_claims,
     
     abstract, claims = input['abstract'], input['claims']
     
-    input_abstract = tokenizer_abstract(abstract, return_tensors='pt')
-    input_claims = tokenizer_claims(claims, return_tensors='pt')
+    encoding_abstract = tokenizer_abstract(abstract, return_tensors='pt', truncation=True, padding='max_length')
+    encoding_claims = tokenizer_claims(claims, return_tensors='pt', truncation=True, padding='max_length')
+    
+    input_abstract = encoding_abstract['input_ids'].to(device)
+    attention_mask_abstract = encoding_abstract['attention_mask'].to(device)
+    input_claims = encoding_claims['input_ids'].to(device)
+    attention_mask_claims = encoding_claims['attention_mask'].to(device)
     
     with torch.no_grad():
-        outputs_abstract = model_abstract(**input_abstract)
-        outputs_claims = model_claims(**input_claims)
+        outputs_abstract = model_abstract(input_ids=input_abstract, attention_mask=attention_mask_abstract)
+        outputs_claims = model_claims(input_ids=input_claims, attention_mask=attention_mask_claims)
         
     combined_prob = (outputs_abstract.logits.softmax(dim=1) + outputs_claims.logits.softmax(dim=1)) / 2
-    label = torch.argmax(combined_prob, dim=1)
+    label = torch.argmax(combined_prob, axis=1).flatten()
     
     return label, combined_prob
     
@@ -53,7 +59,7 @@ if __name__ == '__main__':
     form = st.form('patent-prediction-form')
     dropdown = [example['patent_number'] for example in dataset]
     
-    input_application = form.selectbox('Select a patent\'s application number', patents_dropdown)
+    input_application = form.selectbox('Select a patent\'s application number', dropdown)
     submit = form.form_submit_button("Submit")
     
     if submit:
@@ -62,6 +68,6 @@ if __name__ == '__main__':
         label, prob = predict(model_abstract, model_claims, tokenizer_abstract, tokenizer_claims, input)
         
         st.write(label)
-        st.write(predict)
+        st.write(prob)
         st.write(input['decision'])