Merge branch 'main' of https://huggingface.co/spaces/polygraf-ai/article_writer

app.py

@@ -7,6 +7,18 @@ from ai_generate import generate
 import requests
 from gptzero_free import GPT2PPL
 import language_tool_python
+import re
+from gradio_client import Client
+
+highlighter_client = Client("polygraf-ai/highlighter_test")
+
+
+def remove_bracketed_numbers(text):
+    # This regex pattern matches any occurrence of [{int}]
+    pattern = r"^\[\d+\]"
+    # re.sub() replaces all occurrences of the pattern with an empty string
+    cleaned_text = re.sub(pattern, "", text)
+    return cleaned_text
 
 
 def clean_text(text: str) -> str:
@@ -83,12 +95,8 @@ def regenerate_prompt(settings: Dict[str, str]) -> str:
     
     Comments:
     - {settings['user_comments']}
-
-    Additional requirements:
-    - Include {settings['num_examples']} relevant examples or case studies
-    - Incorporate data or statistics from {', '.join(settings['references'])}
-    - End with a {settings['conclusion_type']} conclusion
-    - Add a "References" section at the end with at least 3 credible sources, formatted as [1], [2], etc.
+    - The originial content should not be changed. Make minor modifications based on user comments above.
+    - Keep the references the same as the given text in the same format.
     - Do not make any headline, title bold.
     
     Ensure proper paragraph breaks for better readability.
@@ -181,13 +189,15 @@ def humanize(
 
 
 ai_check_options = [
-    "Polygraf AI",
+    "Polygraf AI Base",
+    "Polygraf AI 2 Sentence Block",
+    "Polygraf AI 3 Sentence Block",
     # "Sapling AI",
     "GPTZero",
 ]
 
 
-def ai_generated_test_polygraf(text: str) -> Dict:
+def ai_generated_test_polygraf(text: str):
     url = "http://34.66.10.188/ai-vs-human"
     access_key = "6mcemwsFycVVgVjMFwKXki3zJka1r7N4u$Z0Y|x$gecC$hdNtpQf-SpL0+=k;u%BZ"
     headers = {"ACCESS_KEY": access_key}
@@ -207,18 +217,41 @@ def ai_generated_test_gptzero(text):
     gptzero_model = GPT2PPL()
     result = gptzero_model(text)
     print(result)
-    return result
+    return result, None
+
+
+def highlighter_polygraf(text, model="Polygraf AI Base"):
+    if model == "Polygraf AI Base":
+        result = highlighter_client.predict(text=text, model="BC Original", api_name="/process_text")
+        ai_score = result[0]["confidences"][0]["confidence"]
+        ai_res = {"AI": ai_score, "HUMAN": 1 - ai_score}
+        return ai_res, result[1]
+    elif model == "Polygraf AI 2 Sentence Block":
+        result = highlighter_client.predict(text=text, model="BC 2 Sentences", api_name="/process_text")
+        ai_score = result[0]["confidences"][0]["confidence"]
+        ai_res = {"AI": ai_score, "HUMAN": 1 - ai_score}
+        return ai_res, result[1]
+    elif model == "Polygraf AI 3 Sentence Block":
+        result = highlighter_client.predict(text=text, model="BC 3 Sentences", api_name="/process_text")
+        ai_score = result[0]["confidences"][0]["confidence"]
+        ai_res = {"AI": ai_score, "HUMAN": 1 - ai_score}
+        return ai_res, result[1]
+    else:
+        result = highlighter_client.predict(text=text, model="BC Original", api_name="/process_text")
+        ai_score = result[0]["confidences"][0]["confidence"]
+        ai_res = {"AI": ai_score, "HUMAN": 1 - ai_score}
+        return ai_res, result[1]
 
 
-def ai_check(text: str, option: str) -> Dict:
-    if option == "Polygraf AI":
-        return ai_generated_test_polygraf(text)
+def ai_check(text: str, option: str):
+    if option.startswith("Polygraf AI"):
+        return highlighter_polygraf(text, option)
     elif option == "Sapling AI":
         return ai_generated_test_sapling(text)
     elif option == "GPTZero":
         return ai_generated_test_gptzero(text)
     else:
-        return ai_generated_test_polygraf(text)
+        return highlighter_polygraf(text, option)
 
 
 def update_visibility_api(model: str):
@@ -236,7 +269,7 @@ def format_references(text: str) -> str:
     in_references = False
 
     for line in lines:
-        if "references" in line.strip().lower():
+        if line.strip().lower() == "references" or line.strip().lower() == "references:":
             in_references = True
             continue
         if in_references:
@@ -246,8 +279,7 @@ def format_references(text: str) -> str:
 
     formatted_refs = []
     for i, ref in enumerate(references, 1):
-        if f"[{i}]" in ref:
-            ref = ref.replace(f"[{i}]", "")
+        ref = remove_bracketed_numbers(ref)
         formatted_refs.append(f"[{i}] {ref}\n")
 
     return "\n\n".join(article_text) + "\n\nReferences:\n" + "\n".join(formatted_refs)
@@ -459,6 +491,7 @@ def create_interface():
                         ai_check_btn = gr.Button("AI Check")
                     ai_check_result = gr.Label(label="AI Check Result")
 
+                highlighted_text = gr.HTML(label="Sentence Breakdown", visible=False)
                 humanize_btn = gr.Button("Humanize")
                 # humanized_output = gr.Textbox(label="Humanized Article", lines=20, elem_classes=["custom-textbox"])
                 humanized_output = gr.Markdown(label="Humanized Article", value="\n\n\n\n", render=True)
@@ -470,6 +503,13 @@ def create_interface():
             else:
                 return gr.update(visible=False)
 
+        def highlight_visible(text):
+            if text.startswith("Polygraf"):
+                return gr.update(visible=True)
+            else:
+                return gr.update(visible=False)
+
+        ai_detector_dropdown.change(highlight_visible, inputs=ai_detector_dropdown, outputs=highlighted_text)
         output_article.change(become_visible, inputs=output_article, outputs=ai_comments)
         ai_comments.change(become_visible, inputs=output_article, outputs=regenerate_btn)
 
@@ -518,6 +558,7 @@ def create_interface():
             ],
             outputs=[output_article],
         )
+
         regenerate_btn.click(
             fn=generate_and_format,
             inputs=[
@@ -544,7 +585,7 @@ def create_interface():
         ai_check_btn.click(
             fn=ai_check,
             inputs=[output_article, ai_detector_dropdown],
-            outputs=[ai_check_result],
+            outputs=[ai_check_result, highlighted_text],
         )
 
         humanize_btn.click(