File size: 4,409 Bytes
2c0549b 66dbc43 2c0549b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 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=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 |