Update app.py

app.py

@@ -1,18 +1,18 @@
 import gradio as gr
-from transformers import DebertaTokenizer, DebertaForSequenceClassification
+from transformers import DebertaTokenizer, DebertaForSequenceClassification, DistilBertTokenizer, DistilBertForSequenceClassification
 from transformers import pipeline
 import json
 
 
-save_path_abstract = './fine-tuned-deberta'
-model_abstract = DebertaForSequenceClassification.from_pretrained(save_path_abstract)
-tokenizer_abstract = DebertaTokenizer.from_pretrained(save_path_abstract)
+save_path_abstract = './fine-tuned-distillberta'
+model_abstract = DistilBertForSequenceClassification.from_pretrained(save_path_abstract)
+tokenizer_abstract = DistilBertTokenizer.from_pretrained(save_path_abstract)
 
 classifier_abstract = pipeline('text-classification', model=model_abstract, tokenizer=tokenizer_abstract)
 
-save_path_essay = './fine-tuned-deberta'
-model_essay = DebertaForSequenceClassification.from_pretrained(save_path_essay)
-tokenizer_essay = DebertaTokenizer.from_pretrained(save_path_essay)
+save_path_essay = './fine-tuned-distillberta'
+model_essay = DistilBertForSequenceClassification.from_pretrained(save_path_essay)
+tokenizer_essay = DistilBertTokenizer.from_pretrained(save_path_essay)
 
 classifier_essay = pipeline('text-classification', model=model_essay, tokenizer=tokenizer_essay)
 
@@ -36,18 +36,13 @@ def process_result_detection_tab(text):
             'Human Written, Machine Polished': float: the probability that the text is human written and machine polished
     '''
     mapping = {'llm': 'Machine Generated', 'human':'Human Written', 'machine-humanized': 'Machine Written, Machine Humanized', 'machine-polished': 'Human Written, Machine Polished'} 
+    result = classifier_abstract(text)
+    result_r = classifier_essay(text)
 
-    # Initialize scores for all classes
-    final_results = {label: 0.0 for label in mapping.values()}
-
-    # Add scores from classifier_abstract
-    if result['label'] in mapping:
-        final_results[mapping[result['label']]] += 0.5 * result['score']
-
-    # Add scores from classifier_essay
-    if result_r['label'] in mapping:
-        final_results[mapping[result_r['label']]] += 0.5 * result_r['score']
+    labels = [mapping[x['label']] for x in result]
+    scores = list(0.5 * np.array([x['score'] for x in result]) + 0.5 * np.array([x['score'] for x in result_r]))
 
+    final_results = dict(zip(labels, scores))
     print(final_results)
     return final_results