Update app.py

app.py

@@ -3,6 +3,9 @@ from transformers import pipeline, AutoModelForSeq2SeqLM, AutoTokenizer
 import os
 import torch
 import spaces
+from datasets import Dataset
+import time
+import datetime
 
 # Define model paths
 MODEL_PATHS = {
@@ -12,14 +15,25 @@ MODEL_PATHS = {
     "Terjman-Supreme-v2": "BounharAbdelaziz/Terjman-Supreme-v2.0"
 }
 
-# Load environment token
+# Load environment tokens
 TOKEN = os.environ['TOKEN']
 
-# Preload models and tokenizers
+# Dataset configuration
+DATASET_REPO = "BounharAbdelaziz/terjman-v2-live-translations"
+# Number of translations to collect before pushing
+BATCH_SIZE = 10  
+# Time in seconds between pushes (1 hour)
+UPDATE_INTERVAL = 3600
+
+# Initialize dataset tracking
+translations_buffer = []
+last_push_time = time.time()
+
 def preload_models():
+    """ Preload models and tokenizers """
     device = "cuda:0" if torch.cuda.is_available() else "cpu"
     print(f"[INFO] Using device: {device}")
-
+    
     # Load Nano and Large models
     nano_large_models = {}
     for model_name in ["Terjman-Nano-v2", "Terjman-Large-v2"]:
@@ -31,7 +45,7 @@ def preload_models():
             device=device if device.startswith("cuda") else -1
         )
         nano_large_models[model_name] = translator
-
+    
     # Load Ultra and Supreme models
     ultra_supreme_models = {}
     for model_name in ["Terjman-Ultra-v2", "Terjman-Supreme-v2"]:
@@ -47,15 +61,46 @@ def preload_models():
             tgt_lang="ary_Arab"
         )
         ultra_supreme_models[model_name] = translator
-
+    
     return nano_large_models, ultra_supreme_models
 
-# Preload all models
-nano_large_models, ultra_supreme_models = preload_models()
+def push_to_hf_dataset():
+    """ Save translations in HF dataset for monitoring """
+    global translations_buffer, last_push_time
+    
+    if not translations_buffer:
+        return
+    
+    try:
+        print(f"[INFO] Pushing {len(translations_buffer)} translations to Hugging Face dataset...")
+        
+        # Create dataset from buffer
+        ds = Dataset.from_dict({
+            "source_text": [item["source_text"] for item in translations_buffer],
+            "translated_text": [item["translated_text"] for item in translations_buffer],
+            "model_used": [item["model_used"] for item in translations_buffer],
+            "timestamp": [item["timestamp"] for item in translations_buffer]
+        })
+        
+        # Push to hub
+        ds.push_to_hub(
+            DATASET_REPO, 
+            token=TOKEN,
+            split=f"translations_{datetime.datetime.now().strftime('%Y%m%d_%H%M%S')}",
+            private=True,
+        )
+        
+        # Clear buffer and reset timer
+        translations_buffer = []
+        last_push_time = time.time()
+        print("[INFO] Successfully pushed translations to Hugging Face dataset")
+    
+    except Exception as e:
+        print(f"[ERROR] Failed to push dataset to Hugging Face: {str(e)}")
 
-# Translation function for Nano and Large models
 @spaces.GPU
 def translate_nano_large(text, model_name):
+    """ Translation function for Nano and Large models """
     translator = nano_large_models[model_name]
     translated = translator(
         text,
@@ -70,47 +115,94 @@ def translate_nano_large(text, model_name):
     )
     return translated[0]["translation_text"]
 
-# Translation function for Ultra and Supreme models
 @spaces.GPU
 def translate_ultra_supreme(text, model_name):
+    """ Translation function for Ultra and Supreme models """
     translator = ultra_supreme_models[model_name]
     translation = translator(text)[0]['translation_text']
     return translation
 
-# Main translation function
 def translate_text(text, model_choice):
+    """ Main translation function """
+    global translations_buffer, last_push_time
+    
+    # Skip empty text
+    if not text or text.strip() == "":
+        return "Please enter text to translate."
+    
+    # Perform translation
     if model_choice in ["Terjman-Nano-v2", "Terjman-Large-v2"]:
-        return translate_nano_large(text, model_choice)
+        translation = translate_nano_large(text, model_choice)
     elif model_choice in ["Terjman-Ultra-v2", "Terjman-Supreme-v2"]:
-        return translate_ultra_supreme(text, model_choice)
+        translation = translate_ultra_supreme(text, model_choice)
     else:
         return "Invalid model selection."
+    
+    # Add to buffer
+    translations_buffer.append({
+        "source_text": text,
+        "translated_text": translation,
+        "model_used": model_choice,
+        "timestamp": datetime.datetime.now().isoformat()
+    })
+    
+    # Check if it's time to push to HF
+    current_time = time.time()
+    if len(translations_buffer) >= BATCH_SIZE or (current_time - last_push_time) >= UPDATE_INTERVAL:
+        push_to_hf_dataset()
+    
+    return translation
 
-# Gradio app
 def gradio_app():
     with gr.Blocks() as app:
         gr.Markdown("# 🇲🇦 Terjman-v2")
         gr.Markdown("Choose a model and enter the English text you want to translate to Moroccan Darija.")
-
+        
         model_choice = gr.Dropdown(
             label="Select Model",
             choices=["Terjman-Nano-v2", "Terjman-Large-v2", "Terjman-Ultra-v2", "Terjman-Supreme-v2"],
             value="Terjman-Ultra-v2"
         )
+        
         input_text = gr.Textbox(label="Input Text", placeholder="Enter text to translate...", lines=3)
         output_text = gr.Textbox(label="Translated Text", interactive=False, lines=3)
+        
         translate_button = gr.Button("Translate")
-
+        
+        # Status message
+        status = gr.Markdown(f"Translations in buffer: 0")
+        
         # Link input and output
+        def translate_and_update_status(text, model):
+            translation = translate_text(text, model)
+            return translation, f"Translations in buffer: {len(translations_buffer)} (Will push when reaching {BATCH_SIZE} or after {UPDATE_INTERVAL/3600} hours)"
+        
         translate_button.click(
-            fn=translate_text,
+            fn=translate_and_update_status,
             inputs=[input_text, model_choice],
-            outputs=output_text
+            outputs=[output_text, status]
         )
-
+        
+        # Add a manual push button for admins
+        with gr.Accordion("Admin Controls", open=False):
+            push_button = gr.Button("Push Current Buffer to HF Dataset")
+            push_button.click(
+                fn=lambda: (push_to_hf_dataset(), f"Pushed translations to HF. Buffer size: {len(translations_buffer)}"),
+                inputs=[],
+                outputs=[status]
+            )
+    
     return app
 
 # Run the app
 if __name__ == "__main__":
+    # Register shutdown handler to save remaining translations
+    import atexit
+    atexit.register(push_to_hf_dataset)
+
+    # Preload all models
+    nano_large_models, ultra_supreme_models = preload_models()
+    
+    # Launch the app
     app = gradio_app()
     app.launch()