Update app.py

app.py

@@ -9,13 +9,24 @@ from torch.nn.utils.rnn import pad_sequence
 import firebase_admin
 from firebase_admin import credentials, firestore
 
+
 # Define the model architecture
 class CTCTransliterator(nn.Module):
-    def __init__(self, input_dim, hidden_dim, output_dim, num_layers=3, dropout=0.3, upsample_factor=3):
+
+    def __init__(self,
+                 input_dim,
+                 hidden_dim,
+                 output_dim,
+                 num_layers=3,
+                 dropout=0.3,
+                 upsample_factor=3):
         super().__init__()
         self.embed = nn.Embedding(input_dim, hidden_dim, padding_idx=0)
-        self.lstm = nn.LSTM(hidden_dim, hidden_dim, num_layers=num_layers,
-                            bidirectional=True, dropout=dropout)
+        self.lstm = nn.LSTM(hidden_dim,
+                            hidden_dim,
+                            num_layers=num_layers,
+                            bidirectional=True,
+                            dropout=dropout)
         self.layer_norm = nn.LayerNorm(hidden_dim * 2)
         self.dropout = nn.Dropout(dropout)
         self.upsample_factor = upsample_factor
@@ -30,20 +41,26 @@ class CTCTransliterator(nn.Module):
 
         # (seq_len, batch, hidden) → (batch, hidden, seq_len)
         x = x.permute(1, 2, 0)
-        x = F.interpolate(x, scale_factor=self.upsample_factor, mode='linear', align_corners=False)
+        x = F.interpolate(x,
+                          scale_factor=self.upsample_factor,
+                          mode='linear',
+                          align_corners=False)
         # → (batch, hidden, seq_len*upsample_factor)
-        x = x.permute(2, 0, 1)  # back to (seq_len*upsample_factor, batch, hidden)
+        x = x.permute(2, 0,
+                      1)  # back to (seq_len*upsample_factor, batch, hidden)
 
         x = self.fc(x)
         x = x.log_softmax(dim=2)
         return x
 
+
 # Firebase Cache System
 class FirebaseCache:
+
     def __init__(self):
         self.db = None
         self.init_firebase()
-        
+
     def init_firebase(self):
         """Initialize Firebase connection"""
         try:
@@ -53,118 +70,127 @@ class FirebaseCache:
                 if os.getenv('FIREBASE_CREDENTIALS'):
                     # Parse credentials from environment variable
                     import base64
-                    cred_data = json.loads(base64.b64decode(os.getenv('FIREBASE_CREDENTIALS')).decode())
+                    cred_data = json.loads(
+                        base64.b64decode(
+                            os.getenv('FIREBASE_CREDENTIALS')).decode())
                     cred = credentials.Certificate(cred_data)
                 elif os.path.exists('firebase-credentials.json'):
                     # For local development
                     cred = credentials.Certificate('firebase-credentials.json')
                 else:
-                    print("No Firebase credentials found. Using local cache fallback.")
+                    print(
+                        "No Firebase credentials found. Using local cache fallback."
+                    )
                     return
-                
+
                 firebase_admin.initialize_app(cred)
                 self.db = firestore.client()
                 print("Firebase initialized successfully!")
             else:
                 self.db = firestore.client()
-                
+
         except Exception as e:
             print(f"Firebase initialization failed: {e}")
             print("Falling back to local cache mode")
             self.db = None
-    
-    def _create_cache_key(self, input_text):
+
+    def _create_cache_key(self, input_text, direction):
         """Create a safe document key for Firestore"""
         import hashlib
         # Create hash to handle special characters and length limits
-        key = f"{input_text}"
+        key = f"{input_text}_{direction}"
         return hashlib.md5(key.encode()).hexdigest()
-    
-    def get(self, input_text):
+
+    def get(self, input_text, direction):
         """Get cached translation from Firebase"""
         if not self.db:
             return None
-            
+
         try:
-            doc_key = self._create_cache_key(input_text)
+            doc_key = self._create_cache_key(input_text, direction)
             doc = self.db.collection('translations').document(doc_key).get()
-            
+
             if doc.exists:
                 data = doc.to_dict()
                 # Update usage count
                 self.db.collection('translations').document(doc_key).update({
-                    'usage_count': data.get('usage_count', 0) + 1,
-                    'last_used': datetime.now()
+                    'usage_count':
+                    data.get('usage_count', 0) + 1,
+                    'last_used':
+                    datetime.now()
                 })
                 print(f"Cache hit: {input_text}")
                 return data.get('output', '')
-            
+
             return None
-            
+
         except Exception as e:
             print(f"Cache read error: {e}")
             return None
-    
-    def set(self, input_text, output):
+
+    def set(self, input_text, direction, output):
         """Store translation in Firebase"""
         if not self.db:
             return False
-            
+
         try:
-            doc_key = self._create_cache_key(input_text)
+            doc_key = self._create_cache_key(input_text, direction)
             doc_data = {
                 'input': input_text,
+                'direction': direction,
                 'output': output,
                 'corrected_output': '',
                 'timestamp': datetime.now(),
                 'last_used': datetime.now(),
                 'usage_count': 1
             }
-            
+
             self.db.collection('translations').document(doc_key).set(doc_data)
             print(f"Cached: {input_text} → {output}")
             return True
-            
+
         except Exception as e:
             print(f"Cache write error: {e}")
             return False
-    
-    def update_correction(self, input_text, corrected_output):
+
+    def update_correction(self, input_text, direction, corrected_output):
         """Update translation with user correction"""
         if not self.db:
             return False
-            
+
         try:
-            doc_key = self._create_cache_key(input_text)
+            doc_key = self._create_cache_key(input_text, direction)
             self.db.collection('translations').document(doc_key).update({
-                'corrected_output': corrected_output,
-                'correction_timestamp': datetime.now()
+                'corrected_output':
+                corrected_output,
+                'correction_timestamp':
+                datetime.now()
             })
             print(f"Correction saved: {input_text} → {corrected_output}")
             return True
-            
+
         except Exception as e:
             print(f"Correction save error: {e}")
             return False
-    
+
     def get_stats(self):
         """Get cache statistics"""
         if not self.db:
             return "Firebase not connected"
-            
+
         try:
             docs = self.db.collection('translations').get()
             total = len(docs)
-            
+
             corrected = 0
             total_usage = 0
-            
+
             for doc in docs:
                 data = doc.to_dict()
                 if data.get('corrected_output'):
                     corrected += 1
                 total_usage += data.get('usage_count', 0)
-            
+
             return f"""
 Cache Statistics:
 • Total translations: {total}
@@ -172,62 +198,67 @@ Cache Statistics:
 • Total usage count: {total_usage}
 • Average usage: {total_usage/total if total > 0 else 0:.1f} per translation
             """.strip()
-            
+
         except Exception as e:
             return f"Error getting stats: {e}"
 
+
 # Load vocabularies and model
 def load_model_and_vocabs():
     device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-    
+
     # Load vocabularies
     with open('latin_stoi.json', 'r', encoding='utf-8') as f:
         latin_stoi = json.load(f)
     with open('latin_itos.json', 'r', encoding='utf-8') as f:
         latin_itos = json.load(f)
-    
+
     with open('arabic_stoi.json', 'r', encoding='utf-8') as f:
         arabic_stoi = json.load(f)
     with open('arabic_itos.json', 'r', encoding='utf-8') as f:
-        arabic_itos= json.load(f)
-    
+        arabic_itos = json.load(f)
+
     # Initialize model
-    model = CTCTransliterator(
-        len(latin_stoi),
-        256,
-        len(arabic_stoi),
-        num_layers=3,
-        dropout=0.3,
-        upsample_factor=2
-    ).to(device)
-        
+    model = CTCTransliterator(len(latin_stoi),
+                              256,
+                              len(arabic_stoi),
+                              num_layers=3,
+                              dropout=0.3,
+                              upsample_factor=2).to(device)
+
     # Load trained weights
-    model.load_state_dict(torch.load('best_model.pth', map_location=device, weights_only=False))
+    model.load_state_dict(
+        torch.load('best_model.pth', map_location=device, weights_only=False))
     model.eval()
-    
-    blank_id = arabic_stoi.get('<blank>', len(arabic_itos)-1)
+
+    blank_id = arabic_stoi.get('<blank>', len(arabic_itos) - 1)
     return model, latin_stoi, latin_itos, arabic_stoi, arabic_itos, blank_id, device
 
+
 # Load everything at startup
-model, latin_stoi, latin_itos, arabic_stoi, arabic_itos, blank_id, device = load_model_and_vocabs()
+model, latin_stoi, latin_itos, arabic_stoi, arabic_itos, blank_id, device = load_model_and_vocabs(
+)
 firebase_cache = FirebaseCache()
 
+
 def encode_text(text, vocab):
     """Encode text using vocabulary"""
-    return torch.tensor([vocab.get(ch, 0) for ch in text.strip()], dtype=torch.long)
+    return torch.tensor([vocab.get(ch, 0) for ch in text.strip()],
+                        dtype=torch.long)
+
 
 def greedy_decode(log_probs, blank_id, itos, stoi):
     """
     Decode CTC outputs using greedy decoding.
     """
-    eos_id = stoi.get('<eos>', len(stoi)-2)
+    eos_id = stoi.get('<eos>', len(stoi) - 2)
     preds = log_probs.argmax(2).T.cpu().numpy()  # (B, T)
     results = []
     raw_results = []
     print(eos_id, blank_id)
     print(stoi)
     print(type(blank_id))
-    print(stoi.get('<eos>',0))
+    print(stoi.get('<eos>', 0))
     for i, pred in enumerate(preds):
         prev = None
         decoded = []
@@ -239,7 +270,7 @@ def greedy_decode(log_probs, blank_id, itos, stoi):
                 break
             # CTC collapse: skip blanks and repeated characters
             if p != blank_id and p != prev:
-                decoded.append(itos[str(p)])  
+                decoded.append(itos[str(p)])
             prev = p
             raw_result.append(itos[str(p)])
 
@@ -249,110 +280,116 @@ def greedy_decode(log_probs, blank_id, itos, stoi):
 
     return results
 
+
 def transliterate_latin_to_arabic(text):
     """Transliterate Latin script to Arabic script with Firebase caching"""
     if not text.strip():
         return ""
-    
+
     # Check Firebase cache first
     cached_result = firebase_cache.get(text, "Latin → Arabic")
     if cached_result:
         return cached_result
-    
+
     try:
         # Encode input text
         src = encode_text(text, latin_stoi).unsqueeze(1).to(device)
-        
+
         # Generate prediction
         with torch.no_grad():
             out = model(src)
-        
+
         # Decode output
         decoded = greedy_decode(out, blank_id, arabic_itos, arabic_stoi)
         result = decoded[0] if decoded else ""
-        
+
         # Cache the result in Firebase
         firebase_cache.set(text, "Latin → Arabic", result)
-        
+
         return result
-    
+
     except Exception as e:
         return f"Error: {str(e)}"
 
+
 def transliterate_arabic_to_latin(text):
     """Transliterate Arabic script to Latin script (placeholder)"""
     return "Arabic to Latin transliteration not implemented yet."
 
-def transliterate(text):
+
+def transliterate(text, direction):
     """Main transliteration function"""
-    return transliterate_latin_to_arabic(text.lower())
+    if direction == "Latin → Arabic":
+        return transliterate_latin_to_arabic(text.lower())
+    else:
+        return transliterate_arabic_to_latin(text)
 
 
-def save_correction(input_text, corrected_output):
+def save_correction(input_text, direction, corrected_output):
     """Save user correction to Firebase"""
-    if firebase_cache.update_correction(input_text, corrected_output):
+    if firebase_cache.update_correction(input_text, direction,
+                                        corrected_output):
         return "Correction saved to the database! Thank you for improving the model."
     else:
         return "Could not save correction to databse."
 
+
 # Arabic keyboard layout
-arabic_keys = [
-    ['ض', 'ص', 'ث', 'ق', 'ف', 'غ', 'ع', 'ه', 'خ', 'ح', 'ج', 'د'],
-    ['ش', 'س', 'ي', 'ب', 'ل', 'ا', 'ت', 'ن', 'م', 'ك', 'ط'],
-    ['ئ', 'ء', 'ؤ', 'ر', 'لا', 'ى', 'ة', 'و', 'ز', 'ظ'],
-    ['ذ', '١', '٢', '٣', '٤', '٥', '٦', '٧', '٨', '٩', '٠']
-]
+arabic_keys = [['ض', 'ص', 'ث', 'ق', 'ف', 'غ', 'ع', 'ه', 'خ', 'ح', 'ج', 'د'],
+               ['ش', 'س', 'ي', 'ب', 'ل', 'ا', 'ت', 'ن', 'م', 'ك', 'ط'],
+               ['ئ', 'ء', 'ؤ', 'ر', 'لا', 'ى', 'ة', 'و', 'ز', 'ظ'],
+               ['ذ', '١', '٢', '٣', '٤', '٥', '٦', '٧', '٨', '٩', '٠']]
+
 
 # Create Gradio interface
 def create_interface():
-    with gr.Blocks(title="Darija Transliterator", theme=gr.themes.Soft()) as demo:
-        gr.Markdown(
-            """
+    with gr.Blocks(title="Darija Transliterator",
+                   theme=gr.themes.Soft()) as demo:
+        gr.Markdown("""
             # Darija Transliterator
             Convert between Latin script and Arabic script for Moroccan Darija
             
             **Firebase-Powered**: Persistent caching across sessions  
             **Arabic Keyboard**: Built-in Arabic keyboard for corrections  
             **Real-time Stats**: Live usage analytics
-            """
-        )
-        
+            """)
+
         # Stats section
         with gr.Row():
             stats_btn = gr.Button("Show Statistics", variant="secondary")
-            stats_display = gr.Textbox(
-                label="Firebase Statistics", 
-                interactive=False, 
-                visible=False,
-                lines=5
-            )
-        
+            stats_display = gr.Textbox(label="Firebase Statistics",
+                                       interactive=False,
+                                       visible=False,
+                                       lines=5)
+
         with gr.Row():
             with gr.Column(scale=1):
-                
+                direction = gr.Radio(
+                    choices=["Latin → Arabic"],
+                    value="Latin → Arabic",
+                    label="Translation Direction")
+
                 input_text = gr.Textbox(
                     placeholder="Enter text to transliterate...",
                     label="Input Text",
                     lines=4,
-                    max_lines=10
-                )
-                
+                    max_lines=10)
+
                 with gr.Row():
                     clear_btn = gr.Button("Clear", variant="secondary")
-                    translate_btn = gr.Button("Transliterate", variant="primary")
-            
+                    translate_btn = gr.Button("Transliterate",
+                                              variant="primary")
+
             with gr.Column(scale=1):
-                output_text = gr.Textbox(
-                    label="Output",
-                    lines=4,
-                    max_lines=10,
-                    interactive=True
-                )
-                
+                output_text = gr.Textbox(label="Output",
+                                         lines=4,
+                                         max_lines=10,
+                                         interactive=True)
+
                 # Arabic Keyboard
                 gr.Markdown("### Arabic Keyboard")
                 gr.Markdown("*Click letters to edit the output text above*")
-                
+
                 with gr.Group():
                     for row in arabic_keys:
                         with gr.Row():
@@ -360,114 +397,95 @@ def create_interface():
                                 btn = gr.Button(char, size="sm", scale=1)
                                 btn.click(
                                     fn=None,
-                                    js=f"(output_text) => output_text + '{char}'",
+                                    js=
+                                    f"(output_text) => output_text + '{char}'",
                                     inputs=[output_text],
                                     outputs=[output_text],
                                     show_progress=False,
-                                    queue=False
-                                )
-                    
+                                    queue=False)
+
                     with gr.Row():
                         space_btn = gr.Button("Space", size="sm", scale=2)
-                        backspace_btn = gr.Button("⌫ Backspace", size="sm", scale=2)
-                        clear_output_btn = gr.Button("Clear Output", size="sm", scale=2)
-                
+                        backspace_btn = gr.Button("⌫ Backspace",
+                                                  size="sm",
+                                                  scale=2)
+                        clear_output_btn = gr.Button("Clear Output",
+                                                     size="sm",
+                                                     scale=2)
+
                 # Correction system
                 with gr.Group():
                     gr.Markdown("### Correction System")
-                    correction_status = gr.Textbox(
-                        label="Status",
-                        interactive=False,
-                        visible=False
-                    )
-                    save_correction_btn = gr.Button("Save Correction", variant="secondary")
-        
+                    correction_status = gr.Textbox(label="Status",
+                                                   interactive=False,
+                                                   visible=False)
+                    save_correction_btn = gr.Button("Save Correction",
+                                                    variant="secondary")
+
         # Keyboard utility buttons
-        space_btn.click(
-            fn=None,
-            js="(output_text) => output_text + ' '",
-            inputs=[output_text],
-            outputs=[output_text],
-            show_progress=False,
-            queue=False
-        )
-        
-        backspace_btn.click(
-            fn=None,
-            js="(output_text) => output_text.slice(0, -1)",
-            inputs=[output_text],
-            outputs=[output_text],
-            show_progress=False,
-            queue=False
-        )
-        
-        clear_output_btn.click(
-            fn=None,
-            js="() => ''",
-            outputs=[output_text],
-            show_progress=False,
-            queue=False
-        )
-        
+        space_btn.click(fn=None,
+                        js="(output_text) => output_text + ' '",
+                        inputs=[output_text],
+                        outputs=[output_text],
+                        show_progress=False,
+                        queue=False)
+
+        backspace_btn.click(fn=None,
+                            js="(output_text) => output_text.slice(0, -1)",
+                            inputs=[output_text],
+                            outputs=[output_text],
+                            show_progress=False,
+                            queue=False)
+
+        clear_output_btn.click(fn=None,
+                               js="() => ''",
+                               outputs=[output_text],
+                               show_progress=False,
+                               queue=False)
+
         # Stats button
-        stats_btn.click(
-            fn=firebase_cache.get_stats,
-            outputs=[stats_display]
-        ).then(
-            fn=lambda: gr.update(visible=True),
-            outputs=[stats_display]
-        )
-        
+        stats_btn.click(fn=firebase_cache.get_stats,
+                        outputs=[stats_display
+                                 ]).then(fn=lambda: gr.update(visible=True),
+                                         outputs=[stats_display])
+
         # Example inputs
         gr.Markdown("### Examples")
-        examples = [
-            ["makay3nich bli katkhdam bzaf", "Latin → Arabic"],
-            ["rah bayn dkchi li katdir kolchi 3ay9 bik", "Latin → Arabic"],
-            ["wach na9dar nakhod caipirinha, 3afak", "Latin → Arabic"],
-            ["ghadi temchi f lkhedma mzyan", "Latin → Arabic"]
-        ]
-        
-        gr.Examples(
-            examples=examples,
-            inputs=[input_text],
-            outputs=output_text,
-            fn=transliterate,
-            cache_examples=False
-        )
-        
+        examples = [["makay3nich bli katkhdam bzaf", "Latin → Arabic"],
+                    [
+                        "rah bayn dkchi li katdir kolchi 3ay9 bik",
+                        "Latin → Arabic"
+                    ],
+                    ["wach na9dar nakhod caipirinha, 3afak", "Latin → Arabic"],
+                    ["ghadi temchi f lkhedma mzyan", "Latin → Arabic"]]
+
+        gr.Examples(examples=examples,
+                    inputs=[input_text, direction],
+                    outputs=output_text,
+                    fn=transliterate,
+                    cache_examples=False)
+
         # Event handlers
-        translate_btn.click(
-            fn=transliterate,
-            inputs=[input_text],
-            outputs=output_text
-        ).then(
-            fn=lambda: gr.update(visible=True),
-            outputs=[correction_status]
-        )
-        
-        clear_btn.click(
-            fn=lambda: ("", ""),
-            outputs=[input_text, output_text]
-        )
-        
-        input_text.submit(
-            fn=transliterate,
-            inputs=[input_text],
-            outputs=output_text
-        )
-        
-        save_correction_btn.click(
-            fn=save_correction,
-            inputs=[input_text, output_text],
-            outputs=[correction_status]
-        ).then(
-            fn=lambda: gr.update(visible=True),
-            outputs=[correction_status]
-        )
-        
+        translate_btn.click(fn=transliterate,
+                            inputs=[input_text, direction],
+                            outputs=output_text).then(
+                                fn=lambda: gr.update(visible=True),
+                                outputs=[correction_status])
+
+        clear_btn.click(fn=lambda: ("", ""), outputs=[input_text, output_text])
+
+        input_text.submit(fn=transliterate,
+                          inputs=[input_text, direction],
+                          outputs=output_text)
+
+        save_correction_btn.click(fn=save_correction,
+                                  inputs=[input_text, direction, output_text],
+                                  outputs=[correction_status]).then(
+                                      fn=lambda: gr.update(visible=True),
+                                      outputs=[correction_status])
+
         # Information
-        gr.Markdown(
-            """
+        gr.Markdown("""
             ### About
             This model transliterates Moroccan Darija between Latin and Arabic scripts using a CTC-based neural network.
             
@@ -482,12 +500,12 @@ def create_interface():
             1. Use the Arabic keyboard to correct any wrong translations
             2. Click "Save Correction" to store your improvement
             3. Your corrections help train better models for everyone!
-            """
-        )
-    
+            """)
+
     return demo
 
+
 # Launch the app
 if __name__ == "__main__":
     demo = create_interface()
-    demo.launch(share=True)
+    demo.launch(share=True)