import torch
import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer

# Function to load the model and tokenizer (only needs to run once)
def load_model():
    model_id = "microsoft/bitnet-b1.58-2B-4T"
    tokenizer = AutoTokenizer.from_pretrained(model_id)
    model = AutoModelForCausalLM.from_pretrained(
        model_id,
        torch_dtype=torch.bfloat16,
        device_map="auto"  # This will use available GPU if present
    )
    return model, tokenizer

# Load the model and tokenizer
print("Loading model, please wait...")
model, tokenizer = load_model()
print("Model loaded successfully!")

# List of supported languages
SUPPORTED_LANGUAGES = [
    "English", "Spanish", "French", "German", "Chinese", 
    "Japanese", "Russian", "Arabic", "Portuguese", "Italian"
]

def translate_text(input_text, source_lang, target_lang, max_length=4096):
    """
    Translates text from source language to target language using the BitNet model
    """
    if not input_text.strip():
        return "Please enter some text to translate."
    
    # Create a translation prompt
    prompt = f"""Translate the following {source_lang} text to {target_lang}.
    
{source_lang} text: {input_text}

{target_lang} translation:"""

    # Create inputs for the model
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    
    # Generate translation
    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=max_length,
            do_sample=False,  # Use greedy decoding for translation
            temperature=0.1,  # Low temperature for more deterministic output
        )
    
    # Extract only the generated part (the translation)
    translated_text = tokenizer.decode(outputs[0][inputs['input_ids'].shape[1]:], skip_special_tokens=True)
    
    return translated_text.strip()

# Define the Gradio interface
def create_translation_interface():
    with gr.Blocks(title="BitNet Multilingual Translation Tool") as demo:
        gr.Markdown("# 🌍 BitNet Multilingual Translation Tool")
        gr.Markdown("A lightweight translation application powered by Microsoft's BitNet b1.58 2B4T model.")
        
        with gr.Row():
            with gr.Column():
                source_lang = gr.Dropdown(
                    choices=SUPPORTED_LANGUAGES, 
                    value="English", 
                    label="Source Language"
                )
                input_text = gr.Textbox(
                    lines=5, 
                    placeholder="Enter text to translate...",
                    label="Input Text"
                )
                
            with gr.Column():
                target_lang = gr.Dropdown(
                    choices=SUPPORTED_LANGUAGES, 
                    value="Spanish", 
                    label="Target Language"
                )
                output_text = gr.Textbox(
                    lines=5, 
                    label="Translated Text"
                )
        
        translate_btn = gr.Button("Translate")
        translate_btn.click(
            fn=translate_text,
            inputs=[input_text, source_lang, target_lang],
            outputs=output_text
        )
        
        # Add some example inputs
        examples = [
            ["Hello, how are you today?", "English", "Spanish"],
            ["I'd like to learn more about artificial intelligence.", "English", "French"],
            ["The weather is beautiful today.", "English", "German"],
            ["Could you please help me find the nearest restaurant?", "English", "Japanese"],
        ]
        gr.Examples(examples=examples, inputs=[input_text, source_lang, target_lang])
        
        gr.Markdown("""
        ## About
        This application uses Microsoft's BitNet b1.58 2B4T, a 1-bit Large Language Model, for translation tasks.
        The model runs efficiently on consumer hardware due to its 1-bit architecture, offering significant
        advantages in memory usage, energy consumption, and latency.
        
        Note: Translation quality may vary by language pair. This is a demonstration of the lightweight model's capabilities.
        """)
        
    return demo

# Create and launch the Gradio interface
if __name__ == "__main__":
    demo = create_translation_interface()
    demo.launch(share=True)  # Set share=False if you don't want a public link