fix model caching error

app.py

@@ -13,6 +13,9 @@ import psutil
 with open("hit_log.txt", mode='a') as file:
     file.write(str(datetime.datetime.now()) + '\n')
 
+
+MAX_GRAPH_ROWS = 10
+
 MODEL_DESC = {
     'Bart MNLI': """Bart with a classification head trained on MNLI.\n\nSequences are posed as NLI premises and topic labels are turned into premises, i.e. `business` -> `This text is about business.`""",
     'Bart MNLI + Yahoo Answers': """Bart with a classification head trained on MNLI and then further fine-tuned on Yahoo Answers topic classification.\n\nSequences are posed as NLI premises and topic labels are turned into premises, i.e. `business` -> `This text is about business.`""",
@@ -58,10 +61,22 @@ models = load_models()
 def load_tokenizer(tok_id):
     return AutoTokenizer.from_pretrained(tok_id)
 
-@st.cache(allow_output_mutation=True, show_spinner=False)
-def get_most_likely(nli_model_id, sequence, labels, hypothesis_template, multi_class, do_print_code):
-    classifier = pipeline('zero-shot-classification', model=models[nli_model_id], tokenizer=load_tokenizer(nli_model_id), device=device)
-    outputs = classifier(sequence, labels, hypothesis_template, multi_class)
+@st.cache(allow_output_mutation=True, show_spinner=False, hash_funcs={
+    torch.nn.Parameter: lambda _: None
+})
+def get_most_likely(nli_model_id, sequence, labels, hypothesis_template, multi_class):
+    classifier = pipeline(
+        'zero-shot-classification',
+        model=models[nli_model_id],
+        tokenizer=load_tokenizer(nli_model_id),
+        device=device
+    )
+    outputs = classifier(
+        sequence,
+        candidate_labels=labels,
+        hypothesis_template=hypothesis_template,
+        multi_label=multi_class
+    )
     return outputs['labels'], outputs['scores']
 
 def load_examples(model_id):
@@ -88,7 +103,6 @@ def plot_result(top_topics, scores):
     fig.update_traces(texttemplate='%{text:0.1f}%', textposition='outside')
     st.plotly_chart(fig)
 
-        
 
 def main():
     with open("style.css") as f:
@@ -124,18 +138,11 @@ def main():
         st.markdown(CODE_DESC.format(model_id))
 
     with st.spinner('Classifying...'):
-        top_topics, scores = get_most_likely(model_id, sequence, labels, hypothesis_template, multi_class, do_print_code)
-
-    plot_result(top_topics[::-1][-10:], scores[::-1][-10:])
-
-    if "socat" not in [p.name() for p in psutil.process_iter()]:
-        os.system('socat tcp-listen:8000,reuseaddr,fork tcp:localhost:8001 &')
-
-
+        top_topics, scores = get_most_likely(model_id, sequence, labels, hypothesis_template, multi_class)
 
+    plot_result(top_topics[::-1][-MAX_GRAPH_ROWS:], scores[::-1][-MAX_GRAPH_ROWS:])
 
 
 
 if __name__ == '__main__':
     main()
-