Change to t5

app.py

@@ -1,14 +1,15 @@
 import streamlit as st
-from transformers import pipeline
+from transformers import T5Tokenizer, T5ForConditionalGeneration
 from happytransformer import HappyTextToText, TTSettings
 
-# Initialize the spelling correction pipeline
-fix_spelling = pipeline("text2text-generation", model="oliverguhr/spelling-correction-english-base")
-
 # Initialize the grammar correction model
 happy_tt = HappyTextToText("T5", "vennify/t5-base-grammar-correction")
 args = TTSettings(num_beams=5, min_length=1)
 
+# Initialize T5 spelling correction tokenizer and model
+tokenizer = T5Tokenizer.from_pretrained("thaboe01/t5-spelling-corrector")
+model = T5ForConditionalGeneration.from_pretrained("thaboe01/t5-spelling-corrector", device="cuda")
+
 # Function to split text into chunks
 def split_text(text, chunk_size=500):
     chunks = []
@@ -40,7 +41,8 @@ def main():
         for chunk in text_chunks:
             try:
                 # Spelling correction
-                corrected_spelling = fix_spelling(chunk)[0]['generated_text']
+                input_ids = tokenizer(chunk, return_tensors="pt").input_ids.to("cuda")
+                corrected_spelling = tokenizer.decode(model.generate(input_ids)[0])
                 corrected_spelling_chunks.append(corrected_spelling)
                 
                 # Grammar correction