a

app.py

@@ -1,7 +1,7 @@
 import gradio as gr
 
-from typo_check import css, process_input
-from title_gen import generate_title, MODEL_OPTIONS
+from typo_check import css, process_input,MODEL_OPTIONS_TYPO
+from title_gen import generate_title, MODEL_OPTIONS_TITLE
 
 
 # Create Gradio interface using the latest syntax
@@ -11,6 +11,9 @@ with gr.Blocks(theme=gr.themes.Default(), css=css) as demo:
             gr.Markdown("# <center>Dhivehi Typo Correction</center>")
             gr.Markdown("This app uses a fine-tuned T5 model to correct typos in Dhivehi text. Enter text with typos and the model will attempt to fix them.")
             
+            with gr.Row():
+                model_choice = gr.Dropdown(choices=list(MODEL_OPTIONS_TYPO.keys()), value="A3 Model", label="Model")
+
             with gr.Row():
                 input_text = gr.Textbox(
                     lines=1,
@@ -19,7 +22,8 @@ with gr.Blocks(theme=gr.themes.Default(), css=css) as demo:
                     rtl=True,
                     elem_classes="textbox1"
                 )
-            
+
+
             with gr.Row():
                 corrected_text = gr.Textbox(
                     lines=1,
@@ -37,7 +41,7 @@ with gr.Blocks(theme=gr.themes.Default(), css=css) as demo:
             submit_btn = gr.Button("ރަނގަޅު ކޮށްލުމަށް",elem_classes="textbox1")  # "Correct" in Dhivehi
             submit_btn.click(
                 fn=process_input,
-                inputs=input_text,
+                inputs=[input_text,model_choice],
                 outputs=[corrected_text, highlighted_diff]
             )
             
@@ -78,7 +82,7 @@ with gr.Blocks(theme=gr.themes.Default(), css=css) as demo:
             with gr.Row():
                 article_content = gr.Textbox(lines=10, label="Article Content", rtl=True, elem_classes="textbox1")
             with gr.Row():
-                model_choice = gr.Dropdown(choices=list(MODEL_OPTIONS.keys()), value="V6 Model", label="Model")
+                model_choice = gr.Dropdown(choices=list(MODEL_OPTIONS_TITLE.keys()), value="V6 Model", label="Model")
             with gr.Row():
                 seed = gr.Slider(0, 10000, value=42, step=1, label="Random Seed")
                 use_sampling = gr.Checkbox(label="Use Sampling (Creative/Random)", value=False)

title_gen.py

@@ -2,9 +2,10 @@ import random
 import numpy as np
 import torch
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+import spaces
 
 # Available models
-MODEL_OPTIONS = {
+MODEL_OPTIONS_TITLE = {
     "V6 Model": "alakxender/t5-divehi-title-generation-v6",
     "XS Model": "alakxender/t5-dhivehi-title-generation-xs"
 }
@@ -28,6 +29,7 @@ prefix = "2title: "
 max_input_length = 512
 max_target_length = 32
 
+@spaces.GPU()
 def generate_title(content, seed, use_sampling, model_choice):
     random.seed(seed)
     np.random.seed(seed)
@@ -35,7 +37,7 @@ def generate_title(content, seed, use_sampling, model_choice):
     if torch.cuda.is_available():
         torch.cuda.manual_seed_all(seed)
 
-    model_dir = MODEL_OPTIONS[model_choice]
+    model_dir = MODEL_OPTIONS_TITLE[model_choice]
     tokenizer, model = get_model_and_tokenizer(model_dir)
 
     input_text = prefix + content.strip()

typo_check.py

@@ -5,19 +5,25 @@ import difflib
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 import torch
 import gradio as gr
+import spaces
 
-# Load the fine-tuned model and tokenizer
-MODEL_PATH = "alakxender/dhivehi-quick-spell-check-t5"  # Change this to your model path if different
+# Available models
+MODEL_OPTIONS_TYPO = {
+    "A3 Model": "alakxender/t5-dhivehi-typo-corrector-asr",
+    "XS Model": "alakxender/dhivehi-quick-spell-check-t5"
+}
 
 # Function to load model and tokenizer
-def load_model():
+def load_model(model_choice):
     print("Loading model and tokenizer...")
     try:
-        tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
+        selected_model = MODEL_OPTIONS_TYPO[model_choice]
+
+        tokenizer = AutoTokenizer.from_pretrained(selected_model)
         if tokenizer.pad_token is None:
             tokenizer.pad_token = tokenizer.eos_token
         
-        model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_PATH)
+        model = AutoModelForSeq2SeqLM.from_pretrained(selected_model)
         
         # Move model to GPU if available
         device = "cuda" if torch.cuda.is_available() else "cpu"
@@ -67,7 +73,7 @@ def correct_typo(text, model, tokenizer, device):
         return f"Error: {str(e)}"
 
 # Initialize model and tokenizer
-model, tokenizer, device = load_model()
+model, tokenizer, device = load_model("A3 Model")
 
 if model is None:
     print("Failed to load model. Please check your model and tokenizer paths.")
@@ -103,9 +109,10 @@ def highlight_differences(original, corrected):
     return f'<div class="dhivehi-diff">{" ".join(html_parts)}</div>'
 
 # Function to process the input for Gradio
-def process_input(text):
+@spaces.GPU()
+def process_input(text,model_choice):
     if model is None:
-        load_model()
+        load_model(model_choice)
     
     corrected = correct_typo(text, model, tokenizer, device)
     highlighted = highlight_differences(text, corrected)