Update app.py

app.py

@@ -12,22 +12,34 @@ def main():
     # Input text area
     input_text = st.text_area("Enter text to translate", "")
 
-    if st.button("Translate"):
-        # Perform translation
-        translated_text = translate_text(input_text, model, tokenizer)
+    # Translation buttons
+    translate_hindi = st.button("Hindi")
+    translate_tamil = st.button("Tamil")
+    translate_telugu = st.button("Telugu")
+
+    if translate_hindi:
+        translated_text = translate_text(input_text, model, tokenizer, target_lang="hi_IN")
+        st.write("Translated Text (Hindi):")
+        st.write(translated_text)
+
+    if translate_tamil:
+        translated_text = translate_text(input_text, model, tokenizer, target_lang="ta_IN")
+        st.write("Translated Text (Tamil):")
+        st.write(translated_text)
 
-        # Display translated text
-        st.write("Translated Text:")
+    if translate_telugu:
+        translated_text = translate_text(input_text, model, tokenizer, target_lang="te_IN")
+        st.write("Translated Text (Telugu):")
         st.write(translated_text)
 
-def translate_text(input_text, model, tokenizer):
+def translate_text(input_text, model, tokenizer, target_lang):
     # Tokenize input text
     input_ids = tokenizer(input_text, return_tensors="pt").input_ids
     
     # Generate translation
     generated_tokens = model.generate(
         input_ids=input_ids,
-        forced_bos_token_id=tokenizer.lang_code_to_id["hi_IN"]
+        forced_bos_token_id=tokenizer.lang_code_to_id[target_lang]
     )
 
     # Decode translated text
@@ -36,4 +48,4 @@ def translate_text(input_text, model, tokenizer):
     return translated_text
 
 if __name__ == '__main__':
-    main()
+    main()