Enutrof commited on
Commit
59608ee
1 Parent(s): c3023b3

Added Gradio app and requirements

Browse files
Files changed (2) hide show
  1. app.py +39 -0
  2. requirements.txt +3 -0
app.py ADDED
@@ -0,0 +1,39 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ from simpletransformers.seq2seq import Seq2SeqModel
3
+
4
+ # Define the models' paths
5
+ BM_MODEL_PATH = "Enutrof/marian-mt-en-pcm"
6
+ BBGM_MODEL_PATH = "NITHUB-AI/marian-mt-bbc-en-pcm"
7
+
8
+ #Load models
9
+ bm_model = Seq2SeqModel(encoder_decoder_type="marian", encoder_decoder_name=BM_MODEL_PATH, use_cuda=False)
10
+ bbgm_model = Seq2SeqModel(encoder_decoder_type="marian", encoder_decoder_name=BBGM_MODEL_PATH, use_cuda=False)
11
+
12
+ # Dictionary to easily select model
13
+ models = {
14
+ "BM Model": bm_model,
15
+ "BBGM Model": bbgm_model
16
+ }
17
+
18
+ def translate(model_name, source_sentence, num_beams):
19
+ selected_model = models[model_name]
20
+ predictions = selected_model.predict([source_sentence] * 3, num_beams=int(num_beams), num_return_sequences=3)
21
+ return tuple(predictions)
22
+
23
+ # Gradio interface
24
+ interface = gr.Interface(
25
+ fn=translate,
26
+ inputs=[
27
+ gr.Dropdown(choices=["BM Model", "BBGM Model"], label="Model Selection"),
28
+ gr.Textbox(placeholder="Enter English sentence here...", label="Source Sentence"),
29
+ gr.Slider(minimum=1, maximum=10, default=5, step=1, label="Number of Beams"),
30
+ ],
31
+ outputs=[
32
+ gr.Textbox(label="Prediction 1"),
33
+ gr.Textbox(label="Prediction 2"),
34
+ gr.Textbox(label="Prediction 3"),
35
+ ],
36
+ live=True
37
+ )
38
+
39
+ interface.launch()
requirements.txt ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ gradio==3.42.0
2
+ simpletransformers>=0.64.3
3
+ torch==2.0.1