Update app.py

app.py

@@ -1,33 +1,28 @@
-import streamlit as st
+import gradio as gr
 from transformers import AutoTokenizer, AutoModelForCausalLM
-import warnings
 
-warnings.simplefilter("ignore")
-
-def main():
-    tokenizer = AutoTokenizer.from_pretrained("Unbabel/TowerBase-13B-v0.1")
-    model = AutoModelForCausalLM.from_pretrained("Unbabel/TowerBase-13B-v0.1", device="cuda" if st.session_state.use_gpu else "cpu", load_in_4bit=True)
-
-    languages = ["English", "Spanish", "Vietnamese", "French", "Portuguese"]
-
-    st.sidebar.title("Translation App")
-    st.sidebar.write("Choose source and target languages:")
-
-    source_lang_index = st.sidebar.selectbox("Source Language", languages)
-    target_lang_index = st.sidebar.selectbox("Target Language", languages)
-
-    source_lang = languages.index(source_lang_index)
-    target_lang = languages.index(target_lang_index)
-
-    text = st.text_area(f"Enter text in {source_lang_index}", "")
-
-    if st.button("Translate"):
-        input_text = f"{source_lang_index}: {text}\n{target_lang_index}:"
-        inputs = tokenizer(input_text, return_tensors="pt")
-        outputs = model.generate(**inputs, max_new_tokens=20)
-
-        translated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
-        st.write(f"Translation in {target_lang_index}: {translated_text}")
-
-if __name__ == "__main__":
-    main()
+# Load tokenizer and model
+tokenizer = AutoTokenizer.from_pretrained("Unbabel/TowerBase-13B-v0.1")
+model = AutoModelForCausalLM.from_pretrained("Unbabel/TowerBase-13B-v0.1", device="cuda" if gr.utils.is_using_gpu() else "cpu", load_in_4bit=True)
+
+# Define translation function
+def translate_text(source_lang, target_lang, text):
+    input_text = f"{source_lang}: {text}\n{target_lang}:"
+    inputs = tokenizer(input_text, return_tensors="pt")
+    outputs = model.generate(**inputs, max_length=150)
+    translated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return translated_text
+
+# Define interface
+iface = gr.Interface(
+    fn=translate_text,
+    inputs=[
+        gr.inputs.Dropdown(["English", "Spanish", "Vietnamese", "French", "Portuguese"], label="Source Language"),
+        gr.inputs.Dropdown(["English", "Spanish", "Vietnamese", "French", "Portuguese"], label="Target Language"),
+        gr.inputs.Textbox(lines=5, label="Input Text")
+    ],
+    outputs=gr.outputs.Textbox(label="Translated Text")
+)
+
+# Run the interface
+iface.launch(share=True)