Spaces:
Build error
Build error
| 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 |