Update app.py

app.py

@@ -1,13 +1,14 @@
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline, AutoModelForSequenceClassification
 import os
 from huggingface_hub import login
+import torch
 
-# Load Hugging Face token from environment (set this in Space secrets as HF_TOKEN)
+# Authenticate with Hugging Face
 HUGGINGFACE_TOKEN = os.getenv("HF_TOKEN")
 login(token=HUGGINGFACE_TOKEN)
 
-# Load Phi-4 Mini model
+# Load Phi-4 Mini
 phi_model_id = "microsoft/phi-4-mini-instruct"
 phi_tokenizer = AutoTokenizer.from_pretrained(phi_model_id, token=HUGGINGFACE_TOKEN)
 phi_model = AutoModelForCausalLM.from_pretrained(
@@ -16,9 +17,14 @@ phi_model = AutoModelForCausalLM.from_pretrained(
 phi_pipe = pipeline("text-generation", model=phi_model, tokenizer=phi_tokenizer)
 
 # Load T5 for paraphrasing
-t5_pipe = pipeline("text2text-generation", model="t5-base")
+t5_pipe = pipeline("text2text-generation", model="google-t5/t5-base")
 
-# Helper function to chunk large text
+# 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
 def chunk_text(text, max_tokens=300):
     paragraphs = text.split("\n\n")
     chunks, current = [], ""
@@ -32,7 +38,7 @@ def chunk_text(text, max_tokens=300):
         chunks.append(current.strip())
     return chunks
 
-# Generate output for each instruction
+# Phi-based instruction
 def generate_phi_prompt(text, instruction):
     chunks = chunk_text(text)
     outputs = []
@@ -45,7 +51,7 @@ def generate_phi_prompt(text, instruction):
             outputs.append(result.strip())
     return "\n\n".join(outputs)
 
-# Tool functions
+# Functions for each tool
 def fix_grammar(text):
     return generate_phi_prompt(text, "Correct all grammar and punctuation errors in the following text. Provide only the corrected version:")
 
@@ -63,25 +69,60 @@ def paraphrase(text):
         outputs.append(output)
     return "\n\n".join(outputs)
 
+# Upload/download handlers
+def load_file(file_obj):
+    if file_obj is None:
+        return ""
+    return file_obj.read().decode("utf-8")
+
+def save_file(text):
+    path = "/tmp/output.txt"
+    with open(path, "w", encoding="utf-8") as f:
+        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("# ✍️ Grammar Fixer & Writing Assistant")
-    gr.Markdown("Fix grammar, improve tone and fluency, and paraphrase text with intelligent chunking.")
+    gr.Markdown("# ✍️ AI Writing Assistant + Detector")
+    gr.Markdown("Fix grammar, improve tone and fluency, paraphrase text, detect AI content, and upload/download files.")
 
     with gr.Row():
-        input_text = gr.Textbox(lines=12, label="Enter Your Text")
+        file_input = gr.File(label="📂 Upload .txt File", file_types=[".txt"])
+        load_btn = gr.Button("📥 Load Text")
+        input_text = gr.Textbox(lines=12, label="Or Paste Text")
+
+    load_btn.click(fn=load_file, inputs=file_input, outputs=input_text)
 
     with gr.Row():
         btn_grammar = gr.Button("✔️ Fix Grammar")
         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")
 
     output_text = gr.Textbox(lines=12, label="Output")
+    ai_output = gr.Label(label="AI Detection Result")
 
     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)
+
+    gr.Markdown("## 📤 Download Output")
+    download_btn = gr.Button("💾 Download as .txt")
+    download_file = gr.File(label="Click to download", interactive=True)
+    download_btn.click(fn=save_file, inputs=output_text, outputs=download_file)
 
 demo.launch()