updated

app.py

@@ -1,37 +1,36 @@
 import gradio as gr
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
 
-# Load models
+# Load grammar model
 grammar_model_name = "prithivida/grammar_error_correcter_v1"
-summarizer_model_name = "facebook/bart-large-cnn"
-title_model_name = "t5-base"
-
-# Load tokenizer and models
 grammar_tokenizer = AutoTokenizer.from_pretrained(grammar_model_name)
 grammar_model = AutoModelForSeq2SeqLM.from_pretrained(grammar_model_name)
 
+# Load summarizer
+summarizer = pipeline("summarization", model="facebook/bart-large-cnn")
+
+# Load title generation model
+title_model_name = "t5-base"
 title_tokenizer = AutoTokenizer.from_pretrained(title_model_name)
 title_model = AutoModelForSeq2SeqLM.from_pretrained(title_model_name)
 
-summarizer = pipeline("summarization", model=summarizer_model_name)
-
-# Define functions
-def polish_abstract(text):
-    # Grammar correction
-    inputs = grammar_tokenizer.encode("gec: " + text, return_tensors="pt", max_length=512, truncation=True)
-    outputs = grammar_model.generate(inputs, max_length=512, num_beams=5, early_stopping=True)
-    corrected = grammar_tokenizer.decode(outputs[0], skip_special_tokens=True)
+def polish_abstract(raw_text):
+    # 1. Grammar Correction
+    grammar_inputs = grammar_tokenizer.encode("gec: " + raw_text, return_tensors="pt", truncation=True, max_length=512)
+    grammar_outputs = grammar_model.generate(grammar_inputs, max_length=512, num_beams=5)
+    corrected_text = grammar_tokenizer.decode(grammar_outputs[0], skip_special_tokens=True)
 
-    # Summarization
-    summary = summarizer(corrected, max_length=130, min_length=30, do_sample=False)[0]['summary_text']
+    # 2. Summary (Concise)
+    summary_output = summarizer(corrected_text, max_length=120, min_length=30, do_sample=False)
+    summary_text = summary_output[0]['summary_text']
 
-    # Title generation
-    title_input = "generate title: " + corrected
-    title_inputs = title_tokenizer.encode(title_input, return_tensors="pt", max_length=512, truncation=True)
-    title_outputs = title_model.generate(title_inputs, max_length=15, num_beams=5, early_stopping=True)
-    title = title_tokenizer.decode(title_outputs[0], skip_special_tokens=True)
+    # 3. Title Generation
+    title_prompt = "generate title: " + corrected_text
+    title_inputs = title_tokenizer.encode(title_prompt, return_tensors="pt", truncation=True, max_length=512)
+    title_outputs = title_model.generate(title_inputs, max_length=20, num_beams=5)
+    generated_title = title_tokenizer.decode(title_outputs[0], skip_special_tokens=True)
 
-    return corrected, summary, title
+    return corrected_text, summary_text, generated_title
 
 # Gradio Interface
 iface = gr.Interface(
@@ -43,7 +42,7 @@ iface = gr.Interface(
         gr.Textbox(label="📘 Suggested Title"),
     ],
     title="🧠 AI Abstract & Title Polisher",
-    description="Paste your raw abstract. The AI will fix grammar, generate a concise summary, and suggest a title."
+    description="Paste your rough abstract. This AI tool will correct grammar, generate a concise summary, and suggest a scientific title."
 )
 
 iface.launch()