Update app.py

app.py

@@ -19,12 +19,12 @@ phi_pipe = pipeline("text-generation", model=phi_model, tokenizer=phi_tokenizer)
 # Load T5 for paraphrasing
 t5_pipe = pipeline("text2text-generation", model="google-t5/t5-base")
 
-# Load AI Detector
+# Load AI detector
 ai_model_id = "openai-community/roberta-base-openai-detector"
 ai_tokenizer = AutoTokenizer.from_pretrained(ai_model_id)
 ai_model = AutoModelForSequenceClassification.from_pretrained(ai_model_id)
 
-# Text chunking
+# Chunking helper
 def chunk_text(text, max_tokens=300):
     paragraphs = text.split("\n\n")
     chunks, current = [], ""
@@ -38,7 +38,7 @@ def chunk_text(text, max_tokens=300):
         chunks.append(current.strip())
     return chunks
 
-# Phi-based instruction
+# Phi prompt-based generation
 def generate_phi_prompt(text, instruction):
     chunks = chunk_text(text)
     outputs = []
@@ -51,7 +51,36 @@ def generate_phi_prompt(text, instruction):
             outputs.append(result.strip())
     return "\n\n".join(outputs)
 
-# Functions for each tool
+# AI detection logic
+def detect_ai_text(text):
+    inputs = ai_tokenizer(text, return_tensors="pt", truncation=True, padding=True)
+    with torch.no_grad():
+        logits = ai_model(**inputs).logits
+        probs = torch.softmax(logits, dim=1).squeeze()
+    return probs[0].item(), probs[1].item()  # human, ai
+
+# Auto-rewrite if AI > 50%
+def auto_rewrite_for_human(text):
+    rewritten = generate_phi_prompt(
+        text,
+        "Rewrite the following text so that it is indistinguishable from human writing and avoids AI detection. Be natural and fluent:"
+    )
+    human_score, ai_score = detect_ai_text(rewritten)
+    if ai_score < 0.05:
+        return rewritten, f"✅ Rewritten to pass detection. AI Likelihood: {round(ai_score * 100, 2)}%"
+    else:
+        return rewritten, f"⚠️ Rewritten, but AI Likelihood is still high: {round(ai_score * 100, 2)}%"
+
+# Smart detection & rewrite combo
+def check_and_rewrite(text):
+    human_score, ai_score = detect_ai_text(text)
+    if ai_score > 0.5:
+        rewritten, message = auto_rewrite_for_human(text)
+        return rewritten, {"Likely Human": round(human_score, 2), "Likely AI-Generated": round(ai_score, 2)}, message
+    else:
+        return text, {"Likely Human": round(human_score, 2), "Likely AI-Generated": round(ai_score, 2)}, "✅ Text is human-like. No rewrite needed."
+
+# Tool functions
 def fix_grammar(text):
     return generate_phi_prompt(text, "Correct all grammar and punctuation errors in the following text. Provide only the corrected version:")
 
@@ -69,7 +98,7 @@ def paraphrase(text):
         outputs.append(output)
     return "\n\n".join(outputs)
 
-# Upload/download handlers
+# File utilities
 def load_file(file_obj):
     if file_obj is None:
         return ""
@@ -81,21 +110,10 @@ def save_file(text):
         f.write(text)
     return path
 
-# AI Detection function
-def detect_ai_text(text):
-    inputs = ai_tokenizer(text, return_tensors="pt", truncation=True, padding=True)
-    with torch.no_grad():
-        logits = ai_model(**inputs).logits
-        probs = torch.softmax(logits, dim=1).squeeze()
-    return {
-        "Likely Human": round(probs[0].item(), 2),
-        "Likely AI-Generated": round(probs[1].item(), 2)
-    }
-
 # Gradio UI
 with gr.Blocks() as demo:
-    gr.Markdown("# ✍️ AI Writing Assistant + Detector")
-    gr.Markdown("Fix grammar, improve tone and fluency, paraphrase text, detect AI content, and upload/download files.")
+    gr.Markdown("# ✍️ AI Writing Assistant + AI Detector & Auto-Rewriter")
+    gr.Markdown("Fix grammar, improve tone, paraphrase, detect AI content, and automatically rewrite if needed.")
 
     with gr.Row():
         file_input = gr.File(label="📂 Upload .txt File", file_types=[".txt"])
@@ -109,16 +127,17 @@ with gr.Blocks() as demo:
         btn_tone = gr.Button("🎯 Improve Tone")
         btn_fluency = gr.Button("🔄 Improve Fluency")
         btn_paraphrase = gr.Button("🌀 Paraphrase")
-        btn_detect = gr.Button("🕵️ Detect AI vs Human")
+        btn_detect = gr.Button("🕵️ Detect + Rewrite if AI")
 
     output_text = gr.Textbox(lines=12, label="Output")
     ai_output = gr.Label(label="AI Detection Result")
+    rewrite_status = gr.Textbox(label="Status Message", interactive=False)
 
     btn_grammar.click(fn=fix_grammar, inputs=input_text, outputs=output_text)
     btn_tone.click(fn=improve_tone, inputs=input_text, outputs=output_text)
     btn_fluency.click(fn=improve_fluency, inputs=input_text, outputs=output_text)
     btn_paraphrase.click(fn=paraphrase, inputs=input_text, outputs=output_text)
-    btn_detect.click(fn=detect_ai_text, inputs=input_text, outputs=ai_output)
+    btn_detect.click(fn=check_and_rewrite, inputs=input_text, outputs=[output_text, ai_output, rewrite_status])
 
     gr.Markdown("## 📤 Download Output")
     download_btn = gr.Button("💾 Download as .txt")