app.py

import gradio as gr
import os
import sys
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

try:
    generate_file = hf_hub_download(
        repo_id=repo_id,
        filename="generate.py",
        token=os.environ["HF_TOKEN"]
    )
    os.system(f"cp {generate_file} ./generate.py")
except Exception as e:
    print(f"Error downloading files: {e}")

sys.path.append('.')
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 = ''

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;
}

: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()