Rajut commited on
Commit
3a4df42
·
verified ·
1 Parent(s): 498c1f7

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +28 -0
app.py ADDED
@@ -0,0 +1,28 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ from transformers import AutoTokenizer, AutoModelForCausalLM
3
+
4
+ # Load tokenizer and model
5
+ tokenizer = AutoTokenizer.from_pretrained("Unbabel/TowerBase-13B-v0.1")
6
+ model = AutoModelForCausalLM.from_pretrained("Unbabel/TowerBase-13B-v0.1", device="cuda" if gr.utils.is_using_gpu() else "cpu", load_in_4bit=True)
7
+
8
+ # Define translation function
9
+ def translate_text(source_lang, target_lang, text):
10
+ input_text = f"{source_lang}: {text}\n{target_lang}:"
11
+ inputs = tokenizer(input_text, return_tensors="pt")
12
+ outputs = model.generate(**inputs, max_length=150)
13
+ translated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
14
+ return translated_text
15
+
16
+ # Define interface
17
+ iface = gr.Interface(
18
+ fn=translate_text,
19
+ inputs=[
20
+ gr.inputs.Dropdown(["English", "Spanish", "Vietnamese", "French", "Portuguese"], label="Source Language"),
21
+ gr.inputs.Dropdown(["English", "Spanish", "Vietnamese", "French", "Portuguese"], label="Target Language"),
22
+ gr.inputs.Textbox(lines=5, label="Input Text")
23
+ ],
24
+ outputs=gr.outputs.Textbox(label="Translated Text")
25
+ )
26
+
27
+ # Run the interface
28
+ iface.launch(share=True)