Update app.py

app.py

@@ -22,7 +22,6 @@ from unidecode import unidecode
   
 
 nltk.download('punkt')
-tokenizer = AutoTokenizer.from_pretrained('google-bert/bert-base-uncased')
 
 from writing_analysis import (
     normalize,
@@ -189,7 +188,7 @@ text_bc_model_path = "polygraf-ai/v3-bert-3-2m-trun-bc-lighter-spec"
 text_bc_tokenizer = AutoTokenizer.from_pretrained(text_bc_model_path)
 text_bc_model = AutoModelForSequenceClassification.from_pretrained(text_bc_model_path).to(device)
 
-text_mc_model_path = "polygraf-ai/text-detect-mc-bert-base-uncased-mistral-openai-447k-256"
+text_mc_model_path = "polygraf-ai/ai-text-detection-mc-robert-open-ai-detector-v4"
 text_mc_tokenizer = AutoTokenizer.from_pretrained(text_mc_model_path)
 text_mc_model = AutoModelForSequenceClassification.from_pretrained(text_mc_model_path).to(device)
 
@@ -214,11 +213,17 @@ def update_character_count(text):
     return f"{len(text)} characters"
 
 
-def split_text_allow_complete_sentences_nltk(text, max_length=256, tolerance=40, min_last_segment_length=150):
+def split_text_allow_complete_sentences_nltk(text, max_length=256, tolerance=40, min_last_segment_length=150, type_det):
     sentences = nltk.sent_tokenize(text)
     segments = []
     current_segment = []
-    current_length = 0  
+    current_length = 0 
+    
+    if type_det = 'bc':
+        tokenizer = text_bc_tokenizer
+        
+    elif type_det = 'mc':
+        tokenizer = text_mc_tokenizer
     
     for sentence in sentences:
         tokens = tokenizer.tokenize(sentence)
@@ -263,7 +268,7 @@ def split_text_allow_complete_sentences_nltk(text, max_length=256, tolerance=40,
 
 
 def predict_bc(model, tokenizer, text):
-    tokens = tokenizer(
+    tokens = text_bc_tokenizer(
         text, padding='max_length', truncation=True, max_length=256, return_tensors="pt"
     ).to(device)["input_ids"]
     
@@ -273,7 +278,7 @@ def predict_bc(model, tokenizer, text):
     return output_norm
 
 def predict_mc(model, tokenizer, text):
-    tokens = tokenizer(
+    tokens = text_mc_tokenizer(
         text, padding='max_length', truncation=True, return_tensors="pt", max_length=256
     ).to(device)["input_ids"]
     output = model(tokens)
@@ -285,15 +290,19 @@ def ai_generated_test(ai_option, input):
     
     bc_scores = []
     mc_scores = []
-    samples_len = len(split_text_allow_complete_sentences_nltk(input))
-    segments = split_text_allow_complete_sentences_nltk(input)
+    samples_len_bc = len(split_text_allow_complete_sentences_nltk(input, type_det = 'bc'))
+    samples_len_mc = len(split_text_allow_complete_sentences_nltk(input, type_det = 'mc'))
+    segments_bc = split_text_allow_complete_sentences_nltk(input, type_det = 'bc')
+    segments_mc = split_text_allow_complete_sentences_nltk(input, type_det = 'bc')
     
-    for i in range(samples_len):
-        cleaned_text_bc = remove_special_characters(segments[i])
-        cleaned_text_mc = remove_special_characters(segments[i])
+    for i in range(samples_len_bc):
+        cleaned_text_bc = remove_special_characters(segments_bc[i])
         bc_score = predict_bc(text_bc_model, text_bc_tokenizer,cleaned_text_bc )
-        mc_score = predict_mc(text_mc_model, text_mc_tokenizer, cleaned_text_mc)
         bc_scores.append(bc_score)
+        
+    for i in range(samples_len_mc):
+        cleaned_text_mc = remove_special_characters(segments_mc[i])
+        mc_score = predict_mc(text_mc_model, text_mc_tokenizer, cleaned_text_mc)
         mc_scores.append(mc_score)
 
     bc_scores_array = np.array(bc_scores)