Update app.py

app.py

@@ -1,41 +1,36 @@
 # app.py
 
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
 
 model_id = "google/flan-t5-base"  # ~250MB, loads fast
 
-
-# Load tokenizer and model with trust_remote_code
+# Load tokenizer and model correctly for FLAN-T5
 tokenizer = AutoTokenizer.from_pretrained(model_id)
-model = AutoModelForCausalLM.from_pretrained(model_id,
-                                             trust_remote_code=True,
-                                             torch_dtype="auto")
+model = AutoModelForSeq2SeqLM.from_pretrained(model_id)
 
-# Create pipeline (no use_cache to avoid past_key_values crash)
-polisher = pipeline("text-generation",
+# Create pipeline for text2text generation
+polisher = pipeline("text2text-generation",
                     model=model,
-                    tokenizer=tokenizer,
-                    device=-1)
+                    tokenizer=tokenizer)
 
 def oxford_polish_strict(sentence: str) -> str:
     prompt = (
-        "You are an Oxford grammar professor. "
-        "Rewrite the following sentence in formal written English, following the Oxford University Style Guide. "
+        "Rewrite this sentence in formal written English, following the Oxford University Style Guide. "
         "Ensure tense matches time expressions (e.g. 'tomorrow' → future, 'yesterday' → past), "
-        "use British spelling, apply the Oxford comma, and correct uncountable nouns naturally.\n\n"
-        f"Sentence: {sentence}\n\nCorrected:"
+        "use British spelling, apply the Oxford comma, and correct uncountable nouns naturally. "
+        "Sentence: " + sentence
     )
     out = polisher(prompt, max_new_tokens=80, do_sample=False)
-    return out[0]["generated_text"].split("Corrected:")[-1].strip()
+    return out[0]["generated_text"].strip()
 
 # Gradio interface
 demo = gr.Interface(
-    fn=oxford_polish,
+    fn=oxford_polish_strict,
     inputs=gr.Textbox(lines=2, placeholder="Enter a sentence to correct..."),
     outputs=gr.Textbox(label="Oxford-style Correction"),
     title="Oxford Grammar Polisher",
-    description="Rewrite sentences in formal written English using Oxford grammar rules. Powered by Falcon-7B-Instruct."
+    description="Rewrite sentences in formal written English using Oxford grammar rules. Powered by FLAN-T5."
 )
 
 demo.launch()