v6

app.py

@@ -75,10 +75,11 @@ with gr.Blocks(theme=gr.themes.Default(), css=css) as demo:
             
             import sys
             sys.path.append('.')
-            from title_gen import generate_title
-            
+            from title_gen import generate_title,MODEL_OPTIONS
             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")
             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)
@@ -87,7 +88,7 @@ with gr.Blocks(theme=gr.themes.Default(), css=css) as demo:
             generate_btn = gr.Button("Generate Title")
             generate_btn.click(
                 fn=generate_title,
-                inputs=[article_content, seed, use_sampling],
+                inputs=[article_content, seed, use_sampling, model_choice],
                 outputs=generated_title
             )
 
@@ -152,16 +153,9 @@ This is an experimental model trained on a very small dataset of Dhivehi news ar
 - Due to limited training data, the model may not handle all topics or writing styles equally well.
 - Experimental nature means outputs should be reviewed carefully before use.
 
-### Tips for Best Results
-
-- Use well-structured, news-style content similar to the training data
-- Try multiple seeds to get different title options
-- Keep content focused on a single main topic
-- Review generated titles for accuracy and appropriateness
-
 """)
 
 # Launch the app
 if __name__ == "__main__":
-    #demo.launch(server_name="0.0.0.0", server_port=7811)
-    demo.launch()
+    demo.launch(server_name="0.0.0.0", server_port=7811)
+    #demo.launch()

title_gen.py

@@ -3,22 +3,37 @@ import numpy as np
 import torch
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 
-MODEL_DIR = "alakxender/t5-dhivehi-title-generation-xs"
-tokenizer = AutoTokenizer.from_pretrained(MODEL_DIR)
-model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_DIR)
+# Available models
+MODEL_OPTIONS = {
+    "V6 Model": "alakxender/t5-divehi-title-generation-v6",
+    "XS Model": "alakxender/t5-dhivehi-title-generation-xs"
+}
+
+# Cache for loaded models/tokenizers
+MODEL_CACHE = {}
+
+def get_model_and_tokenizer(model_dir):
+    if model_dir not in MODEL_CACHE:
+        tokenizer = AutoTokenizer.from_pretrained(model_dir)
+        model = AutoModelForSeq2SeqLM.from_pretrained(model_dir)
+        MODEL_CACHE[model_dir] = (tokenizer, model)
+    return MODEL_CACHE[model_dir]
 
 prefix = "2title: "
 
 max_input_length = 512
 max_target_length = 32
 
-def generate_title(content, seed, use_sampling):
+def generate_title(content, seed, use_sampling, model_choice):
     random.seed(seed)
     np.random.seed(seed)
     torch.manual_seed(seed)
     if torch.cuda.is_available():
         torch.cuda.manual_seed_all(seed)
 
+    model_dir = MODEL_OPTIONS[model_choice]
+    tokenizer, model = get_model_and_tokenizer(model_dir)
+
     input_text = prefix + content.strip()
     inputs = tokenizer(
         input_text,