Update demo/demo.py

demo/demo.py

@@ -1,43 +1,56 @@
 import gradio as gr
 from transformers import AutoTokenizer, AutoModel
-# from MMD_calculate import mmd_two_sample_baseline  # Adjust path based on your structure
-# from utils_MMD import extract_features  # Example helper from your utils
+from utils_MMD import extract_features  # Adjust the import path
+from MMD_calculate import mmd_two_sample_baseline  # Adjust the import path
 
 MINIMUM_TOKENS = 64
+THRESHOLD = 0.5  # Threshold for classification
 
 def count_tokens(text, tokenizer):
+    """
+    Counts the number of tokens in the text using the provided tokenizer.
+    """
     return len(tokenizer(text).input_ids)
 
-def run_test_power(model_name, tokenizer_name, real_text, generated_text, N):
+def run_test_power(model_name, real_text, generated_text, N=10):
     """
     Runs the test power calculation for provided real and generated texts.
-    """
 
-    # load tokenizer and model
-    tokenizer = AutoTokenizer.from_pretrained(model_name).cuda()
-    model = AutoModel.from_pretrained(model)
+    Args:
+        model_name (str): Hugging Face model name.
+        real_text (str): Example real text for comparison.
+        generated_text (str): The input text to classify.
+        N (int): Number of repetitions for MMD calculation.
 
+    Returns:
+        str: "Prediction: Human" or "Prediction: AI".
+    """
+    # Load tokenizer and model
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    model = AutoModel.from_pretrained(model_name).cuda()
+    model.eval()
+
+    # Ensure minimum token length
     if count_tokens(real_text, tokenizer) < MINIMUM_TOKENS or count_tokens(generated_text, tokenizer) < MINIMUM_TOKENS:
-        return "Too short length. Need minimum 64 tokens to calculated Test Power."
+        return "Too short length. Need a minimum of 64 tokens to calculate Test Power."
 
     # Extract features
-    fea_real_ls = extract_features(model_name, tokenizer_name, [real_text])
-    fea_generated_ls = extract_features(model_name, tokenizer_name, [generated_text])
+    fea_real_ls = extract_features([real_text], tokenizer, model)
+    fea_generated_ls = extract_features([generated_text], tokenizer, model)
 
-    #  Calculate test power list
-    test_power_ls = mmd_two_sample_baseline(fea_real_ls, fea_generated_ls, N=10)
+    # Calculate test power list
+    test_power_ls = mmd_two_sample_baseline(fea_real_ls, fea_generated_ls, N=N)
 
     # Compute the average test power value
     power_test_value = sum(test_power_ls) / len(test_power_ls)
 
     # Classify the text
-    if power_test_value < threshold:
+    if power_test_value < THRESHOLD:
         return "Prediction: Human"
     else:
         return "Prediction: AI"
 
-
-
+# CSS for custom styling
 css = """
 #header { text-align: center; font-size: 1.5em; margin-bottom: 20px; }
 #output-text { font-weight: bold; font-size: 1.2em; }
@@ -78,9 +91,9 @@ with gr.Blocks(css=css) as app:
         clear_button = gr.Button("Clear", variant="secondary")
     with gr.Row():
         output = gr.Textbox(
-            label = "Prediction",
-            placeholder = "Prediction: Human or AI",
-            elem_id = "output-text",
+            label="Prediction",
+            placeholder="Prediction: Human or AI",
+            elem_id="output-text",
         )
     with gr.Accordion("Disclaimer", open=False):
         gr.Markdown(
@@ -102,7 +115,9 @@ with gr.Blocks(css=css) as app:
             ```
             """
         )
-    submit_button.click(detect_text, inputs=[input_text, model_name], outputs=output)
+    submit_button.click(
+        run_test_power, inputs=[model_name, "The cat sat on the mat.", input_text], outputs=output
+    )
     clear_button.click(lambda: ("", ""), inputs=[], outputs=[input_text, output])
 
 app.launch()