Update app.py

app.py

@@ -1,23 +1,42 @@
-import gradio as gr
-from transformers import pipeline
+import streamlit as st
+import torch
+import os
 
-# Load your model from the Hugging Face repository
-model_name = "SLPG/English_to_Urdu_Unsupervised_MT"
-translator = pipeline("translation_en_to_ur", model=model_name)
+# Load the model checkpoint
+model_path = "https://huggingface.co/SLPG/English_to_Urdu_Unsupervised_MT/tree/main/checkpoint_8_96000.pt"
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-# Define the translation function
-def translate_text(input_text):
-    result = translator(input_text, max_length=400)
-    return result[0]['translation_text']
+# Define a function to load the model
+def load_model(model_path):
+    model = torch.load(model_path, map_location=device)
+    model.eval()
+    return model
 
-# Create the Gradio interface
-interface = gr.Interface(
-    fn=translate_text,
-    inputs=gr.inputs.Textbox(lines=2, placeholder="Enter text to translate"),
-    outputs="text",
-    title="Translation Model Inference",
-    description="Translate text using the Hugging Face translation model."
-)
+# Load the dictionaries
+def load_dictionary(dict_path):
+    with open(dict_path, 'r') as file:
+        dictionary = {line.split()[0]: i for i, line in enumerate(file.readlines())}
+    return dictionary
 
-# Launch the interface
-interface.launch()
+# Translation function
+def translate(model, input_text, src_dict, tgt_dict):
+    # Implement the logic to translate using your model
+    # This is a placeholder, modify according to your model's requirements
+    translated_text = "Translated text here"
+    return translated_text
+
+# Load model and dictionaries
+model = load_model(model_path)
+src_dict = load_dictionary("SLPG/English_to_Urdu_Unsupervised_MT/dict.en.txt")
+tgt_dict = load_dictionary("SLPG/English_to_Urdu_Unsupervised_MT/dict.ur.txt")
+
+# Streamlit interface
+st.title("Translation Model Inference")
+input_text = st.text_area("Enter text to translate", "")
+
+if st.button("Translate"):
+    if input_text:
+        translated_text = translate(model, input_text, src_dict, tgt_dict)
+        st.write(f"Translated Text: {translated_text}")
+    else:
+        st.write("Please enter text to translate.")