Create app.py

app.py

@@ -0,0 +1,118 @@
+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=150):
+    """
+    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