Test multi lingual function

app.py

@@ -1,7 +1,42 @@
+import os
+from collections import OrderedDict
+os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"
+
 import gradio as gr
 from shitsu import ShitsuScorer
+from huggingface_hub import hf_hub_download
+
+class OptimizedShitsuScorer:
+    def __init__(self, max_models=2):
+        self.scorers = OrderedDict()
+        self.max_models = max_models
+        self.current_language = None
+
+    def get_scorer(self, language):
+        if language in self.scorers:
+            # Move the accessed language to the end (most recently used)
+            self.scorers.move_to_end(language)
+        else:
+            # If we're at capacity, remove the least recently used model
+            if len(self.scorers) >= self.max_models:
+                self.scorers.popitem(last=False)
+            
+            # Load the new model
+            self.scorers[language] = ShitsuScorer(language)
+        
+        self.current_language = language
+        return self.scorers[language]
+
+    def score(self, text, language):
+        scorer = self.get_scorer(language)
+        return scorer.score([text])[0]
+
+    def get_loaded_languages(self):
+        return list(self.scorers.keys())
 
-scorer = ShitsuScorer('en')
+optimized_scorer = OptimizedShitsuScorer(max_models=2)
+# Preload English model
+optimized_scorer.get_scorer('en')
 
 example_inputs = [
     "The Beatles were a popular band in the 1960s. They released many hit songs.",
@@ -11,13 +46,14 @@ example_inputs = [
     "Can you believe it's already September? This year is flying by!"
 ]
 
-def get_score(user_text):
-    score = scorer.score([user_text])[0]
+def get_score(user_text, language):
+    score = optimized_scorer.score(user_text, language)
     formatted_score = f"{score:.4g}"
-    yield f'<div class="nice-box"> Score: {formatted_score}</div>'
+    loaded_languages = optimized_scorer.get_loaded_languages()
+    return f'<div class="nice-box"> Score: {formatted_score}</div>', f"Currently loaded languages: {', '.join(loaded_languages)}"
+
+language_options = ['am', 'ar', 'bg', 'bn', 'cs', 'da', 'de', 'el', 'en', 'es', 'fa', 'fi', 'fr', 'gu', 'ha', 'hi', 'hu', 'id', 'it', 'ja', 'jv', 'kn', 'ko', 'lt', 'mr', 'nl', 'no', 'yo', 'zh']
 
-#language_options = ['am', 'ar', 'bg', 'bn', 'cs', 'da', 'de', 'el', 'en', 'es', 'fa', 'fi', 'fr', 'gu', 'ha', 'hi', 'hu', 'id', 'it', 'ja', 'jv', 'kn', 'ko', 'lt', 'mr', 'nl', 'no', 'yo', 'zh']
-#https://huggingface.co/spaces/Dusduo/shitsu-text-scorer-demo/shitsu-logo.jpeg
 css = '''
 #gen_btn{height: 100%}
 #title{text-align: center}
@@ -40,26 +76,32 @@ css = '''
     justify-content: center;
 }
 '''
+
 theme = gr.themes.Soft(
     primary_hue="blue",
     secondary_hue="sky",
 )
 
-#"https://huggingface.co/spaces/Dusduo/shitsu-text-scorer-demo/shitsu-logo.jpeg"
 with gr.Blocks(theme=theme, css=css) as demo:
-
     title = gr.HTML(
         """<h1><img src="https://huggingface.co/spaces/Dusduo/shitsu-text-scorer-demo/resolve/main/shitsu-logo.jpeg" alt="LightBlue"> Shitsu Text Scorer</h1>""",
         elem_id="title",
     )
     gr.Markdown(
-    """This is a demo of [Shitsu text scorer](https://huggingface.co/lightblue/shitsu_text_scorer) for English, which scores text based on the amount of useful, textbook-like information in it.
+    """This is a demo of [Shitsu text scorer](https://huggingface.co/lightblue/shitsu_text_scorer) for multiple languages, which scores text based on the amount of useful, textbook-like information in it.
     
     It outputs a score generally between 0 and 1 but can exceed both of these bounds as it is a regressor.
     """
     )
     with gr.Row():
         user_text = gr.Textbox(label='Input text', placeholder='Type something here...')
+        language_choice = gr.Dropdown(
+            choices=language_options,
+            label="Choose a language",
+            info="Type to search",
+            value="en",
+            allow_custom_value=True,
+        )
         with gr.Column(scale=0):
             submit_btn = gr.Button("Submit")
             score = gr.HTML(
@@ -67,16 +109,9 @@ with gr.Blocks(theme=theme, css=css) as demo:
                 label="Output"
             )
     
-    # Add examples to the textbox
-    gr.Examples(examples=example_inputs, inputs=user_text)
+    loaded_languages = gr.Markdown("Currently loaded languages: en")
     
-    # Searchable dropdown to choose language   
-    # language_choice = gr.Dropdown(
-    #     choices=language_options,
-    #     label="Choose a language",
-    #     info="Type to search",
-    #     allow_custom_value=True,  # Allows typing to search
-    # )
+    gr.Examples(examples=example_inputs, inputs=user_text)
     
     gr.Markdown(
     """
@@ -88,8 +123,6 @@ with gr.Blocks(theme=theme, css=css) as demo:
     """
     )
     
-    # Set up the user submit events
-    user_text.submit(get_score, inputs=[user_text], outputs=[score])
-    submit_btn.click(get_score, inputs=[user_text], outputs=[score])
+    submit_btn.click(get_score, inputs=[user_text, language_choice], outputs=[score, loaded_languages])
     
 demo.launch()

requirements.txt

@@ -1 +1,3 @@
-git+https://github.com/lightblue-tech/shitsu.git
+git+https://github.com/lightblue-tech/shitsu.git
+hf-transfer
+huggingface_hub[hf_transfer]