Update app.py

app.py

@@ -1,14 +1,21 @@
 import streamlit as st
 import torch
+from fairseq.models.transformer import TransformerModel
 import os
 import requests
 
-# Define the URL of your model file
-model_url = "https://huggingface.co/SLPG/English_to_Urdu_Unsupervised_MT/resolve/main/checkpoint_8_96000.pt"
-model_path = "checkpoint_8_96000.pt"
+# Define the URLs of your model and dictionary files
+model_url = "https://huggingface.co/SLPG/English_to_Urdu_Unsupervised_MT/resolve/main/sent_iwslt-bt-enur_42.pt"
+dict_en_url = "https://huggingface.co/SLPG/English_to_Urdu_Unsupervised_MT/resolve/main/dict.en.txt"
+dict_ur_url = "https://huggingface.co/SLPG/English_to_Urdu_Unsupervised_MT/resolve/main/dict.ur.txt"
 
-# Define a function to download the model file
-def download_model(url, file_path):
+# Define the paths to save the downloaded files
+model_path = "sent_iwslt-bt-enur_42.pt"
+dict_en_path = "dict.en.txt"
+dict_ur_path = "dict.ur.txt"
+
+# Define a function to download files
+def download_file(url, file_path):
     if not os.path.exists(file_path):
         with requests.get(url, stream=True) as r:
             r.raise_for_status()
@@ -17,35 +24,17 @@ def download_model(url, file_path):
                     f.write(chunk)
     return file_path
 
-# Load the model checkpoint
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
-# Define a function to load the model
-def load_model(model_path):
-    model = torch.load(model_path, map_location=device)
-    model.eval()
-    return 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
-
-# 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
-
-# Download the model file
-download_model(model_url, model_path)
+# Download the model and dictionary files
+download_file(model_url, model_path)
+download_file(dict_en_url, dict_en_path)
+download_file(dict_ur_url, dict_ur_path)
 
-# Load model and dictionaries
-model = load_model(model_path)
-src_dict = load_dictionary("path/to/dict.en.txt")
-tgt_dict = load_dictionary("path/to/dict.ur.txt")
+# Load the model
+en_ur_model = TransformerModel.from_pretrained(
+    '.',
+    checkpoint_file=model_path,
+    data_name_or_path='.'
+)
 
 # Streamlit interface
 st.title("Translation Model Inference")
@@ -53,7 +42,7 @@ 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}")
+        output_text = en_ur_model.translate(input_text)
+        st.write(f"Translated Text: {output_text}")
     else:
         st.write("Please enter text to translate.")