Update app.py

app.py

@@ -1,15 +1,15 @@
 import gradio as gr
 import os
-import torch
-import threading
-from huggingface_hub import hf_hub_download, login
 import sys
-import gc
+import base64
+import threading
+from huggingface_hub import login, hf_hub_download
 
 login(token=os.environ["HF_TOKEN"])
 tokenizer = os.environ["TOKENIZER"]
 repo_id = os.environ["REPO_ID"]
 
+import torch
 torch.set_num_threads(1)
 if torch.cuda.is_available():
     torch.backends.cudnn.benchmark = True
@@ -17,64 +17,194 @@ if torch.cuda.is_available():
 try:
     generate_file = hf_hub_download(
         repo_id=repo_id,
-        filename="gen.py",
+        filename="generate.py",
         token=os.environ["HF_TOKEN"]
     )
-    os.system(f"cp {generate_file} ./gen.py")
-    
+    os.system(f"cp {generate_file} ./generate.py")
 except Exception as e:
     print(f"Error downloading files: {e}")
-    model_path = None
 
 sys.path.append('.')
-from gen import chat_interface
+from generate import chat_interface, init_model
 
+LOGO_PATH = "static/logo.png"
+if os.path.isfile(LOGO_PATH):
+    with open(LOGO_PATH, "rb") as f:
+        LOGO_B64 = base64.b64encode(f.read()).decode()
+    LOGO_HTML = f'<img src="data:image/png;base64,{LOGO_B64}" alt="Maruth Labs Logo" style="height:40px;">'
+    FAVICON_HTML = f'<link rel="icon" type="image/png" href="data:image/png;base64,{LOGO_B64}">'
+else:
+    LOGO_HTML = '<div style="width:40px;height:40px;background:#ccc;border-radius:4px;"></div>'
+    FAVICON_HTML = ''
 
-with gr.Blocks(title="Madhuram Model Chat") as demo:
-    with gr.Column(elem_id="main-container"):
-        gr.Markdown("# Madhuram Chat Interface", elem_classes="center-text")
-        
-        chatbot = gr.ChatInterface(
-            fn=chat_interface,
-            type="messages",
-            title="Madhuram AI Assistant",
-            description="Chat with the Madhuram language model",
-            theme="soft",
-            retry_btn=True,
-            undo_btn=True,
-            clear_btn=True,
-        )
-        
-        gr.Markdown(
-            "*Disclaimer - This is a demo version of Madhuram. It may occasionally "
-            "generate incorrect or incomplete responses. Please verify important "
-            "information independently. The complete model will be available through "
-            "our own playground where the missing features will be incorporated.*", 
-            elem_classes="disclaimer"
-        )
-
-demo.css = """
-#main-container {
-    max-width: 900px;
-    margin: 0 auto;
-    padding: 20px;
-}
-.center-text {
-    text-align: center;
-    margin-bottom: 30px;
+def init_translation_model():
+    success = init_model(tokenizer_path=tokenizer)
+    if success:
+        print("Model loaded successfully!")
+    else:
+        print("Failed to load model")
+
+def translate_text(source_text, source_lang, target_lang, temperature, top_k, repetition_penalty, max_tokens):
+    return chat_interface(
+        source_text, source_lang, target_lang, 
+        temperature, top_k, repetition_penalty, max_tokens
+    )
+
+languages = ["English", "Hindi", "Bengali", "Tamil", "Telugu", "Kannada", "Panjabi"]
+
+css_path = "static/style.css"
+custom_css = open(css_path, encoding="utf-8").read() if os.path.isfile(css_path) else ""
+
+theme_lock_css = """
+.gradio-container .theme-toggle,
+.gradio-container button[aria-label*="theme"],
+.gradio-container button[title*="theme"],
+.gradio-container .settings button,
+.gradio-container [data-testid="theme-toggle"] {
+    display: none !important;
 }
-.disclaimer {
-    text-align: center;
-    color: #666;
-    font-size: 0.9em;
-    margin-top: 30px;
-    padding: 15px;
-    background-color: #f8f9fa;
-    border-radius: 8px;
-    border-left: 4px solid #007bff;
+
+:root { color-scheme: dark !important; }
+
+body, .gradio-container {
+    background-color: #0a1628 !important;
+    color: #e6eef8 !important;
 }
 """
 
+combined_css = custom_css + theme_lock_css
+
+locked_theme = gr.themes.Monochrome(
+    primary_hue="blue",
+    secondary_hue="slate",
+    neutral_hue="slate"
+).set(
+    background_fill_primary="#0a1628",
+    background_fill_secondary="#1f2937", 
+    block_background_fill="#374151",
+    border_color_primary="#374151",
+    color_accent_soft="#2563eb",
+    block_title_text_color="#e6eef8",
+    block_label_text_color="#e6eef8",
+    body_text_color="#e6eef8"
+)
+
+with gr.Blocks(
+    title="Madhuram Translation - MaruthLabs",
+    css=combined_css,
+    theme=locked_theme,
+    js=f"""
+    function() {{
+        {f'document.head.insertAdjacentHTML("beforeend", `{FAVICON_HTML}`);' if FAVICON_HTML else ''}
+        
+        document.documentElement.setAttribute('data-theme', 'dark');
+        document.body.classList.add('dark');
+        document.body.classList.remove('light');
+
+        const observer = new MutationObserver(function(mutations) {{
+            mutations.forEach(function(mutation) {{
+                mutation.addedNodes.forEach(function(node) {{
+                    if (node.nodeType === 1) {{
+                        const toggles = node.querySelectorAll('.theme-toggle, button[aria-label*="theme"], button[title*="theme"]');
+                        toggles.forEach(toggle => toggle.style.display = 'none');
+                    }}
+                }});
+            }});
+        }});
+        observer.observe(document.body, {{ childList: true, subtree: true }});
+    }}
+    """
+) as demo:
+
+    with gr.Row(elem_classes="main-header"):
+        with gr.Column():
+            gr.HTML(f"""
+            <div style="display:flex;align-items:center;justify-content:space-between;width:100%;">
+                <!-- left: logo + text on one line -->
+                <div style="display:flex;align-items:center;">
+                    {LOGO_HTML}
+                    <h3 style="margin-left:8px;margin-top:0;margin-bottom:0;">Maruth Labs</h3>
+                </div>
+
+                <!-- center title -->
+                <div class="main-title"><h1>Madhuram Translation Model</h1></div>
+
+                <!-- spacer to balance flex -->
+                <div style="width:120px;"></div>
+            </div>
+            """)
+
+    with gr.Row(equal_height=False):
+        with gr.Column(scale=1.5, elem_classes="settings-panel"):
+            gr.Markdown("## Translation Settings")
+            with gr.Row():
+                source_lang = gr.Dropdown(choices=languages, label="Source Language", value="English")
+                target_lang = gr.Dropdown(choices=languages, label="Target Language", value="Hindi")
+            swap_btn = gr.Button("Swap Languages", variant="secondary", size="sm")
+
+            with gr.Accordion("Advanced Settings", open=False):
+                temperature = gr.Slider(0.001, 1.001, 0.001, step=0.1, label="Temperature")
+                top_k = gr.Slider(1, 100, 10, step=1, label="Top-k")
+                repetition_penalty = gr.Slider(1.0, 2.0, 1.2, step=0.1, label="Repetition Penalty")
+                max_tokens = gr.Slider(100, 2000, 400, step=50, label="Max Tokens")
+
+        with gr.Column(scale=2, elem_classes="translation-card"):
+            gr.Markdown("## Translation Interface")
+            source_text = gr.Textbox(
+                label="Enter text to translate", 
+                placeholder="Type or paste your text here", 
+                lines=6, 
+                max_lines=12
+            )
+            with gr.Row():
+                translate_btn = gr.Button("Translate", variant="primary", size="lg")
+                clear_btn = gr.Button("Clear All", variant="secondary", size="lg")
+            translated_text = gr.Textbox(
+                label="Translation Result", 
+                lines=6, 
+                max_lines=12, 
+                interactive=False, 
+                placeholder="Translation will appear here"
+            )
+
+    with gr.Row():
+        with gr.Column():
+            gr.Markdown("### Quick Examples")
+            gr.Examples(
+                examples=[
+                    ["Hello, how are you today?", "English", "Hindi"],
+                    ["তুমি কোথায় যাচ্ছ?", "Bengali", "English"],
+                    ["நீங்கள் எப்படி இருக்கிறீர்கள்?", "Tamil", "Telugu"], 
+                    ["ನಿನ್ನ ಹೆಸರು ಏನು?", "Kannada", "English"],
+                    ["ਸਤ ਸ੍ਰੀ ਅਕਾਲ", "Panjabi", "Hindi"],   
+                ],
+                inputs=[source_text, source_lang, target_lang],
+            )
+
+    def swap_languages(src, tgt):
+        return tgt, src
+    
+    def clear_all():
+        return "", ""
+
+    swap_btn.click(
+        fn=swap_languages, 
+        inputs=[source_lang, target_lang], 
+        outputs=[source_lang, target_lang]
+    )
+    
+    clear_btn.click(
+        fn=clear_all, 
+        outputs=[source_text, translated_text]
+    )
+    
+    translate_btn.click(
+        fn=translate_text, 
+        inputs=[source_text, source_lang, target_lang, temperature, top_k, repetition_penalty, max_tokens], 
+        outputs=[translated_text]
+    )
+    
+    demo.load(fn=init_translation_model)
 
 if __name__ == "__main__":
     demo.launch()