Mohamed Aymane Farhi commited on
Commit
9c8ab75
1 Parent(s): 82d8458

Add other models

Browse files
Files changed (1) hide show
  1. app.py +22 -7
app.py CHANGED
@@ -1,16 +1,31 @@
1
  import gradio as gr
2
  from transformers import pipeline
3
 
 
 
 
 
 
 
 
 
 
4
 
5
- pipe = pipeline("translation", model="Helsinki-NLP/opus-mt-en-ber")
6
 
7
- def predict(text):
8
- return pipe(text)[0]["translation_text"]
 
 
 
9
 
10
  demo = gr.Interface(
11
- fn=predict,
12
- inputs='text',
13
- outputs='text',
 
 
 
14
  )
15
 
16
- demo.launch()
 
1
  import gradio as gr
2
  from transformers import pipeline
3
 
4
+ models = [
5
+ "Helsinki-NLP/opus-mt-en-ber",
6
+ "Helsinki-NLP/opus-mt-ber-en",
7
+ "Helsinki-NLP/opus-mt-fr-ber",
8
+ "Helsinki-NLP/opus-mt-ber-fr",
9
+ "Helsinki-NLP/opus-mt-es-ber",
10
+ "Helsinki-NLP/opus-mt-ber-es",
11
+ "Helsinki-NLP/opus-mt-kab-en"
12
+ ]
13
 
14
+ pipes = {}
15
 
16
+ def predict(text, model):
17
+ if model not in pipes:
18
+ pipes[model] = pipeline("translation", model=model)
19
+ pipe = pipes[model]
20
+ return pipe(text)[0]['translation_text']
21
 
22
  demo = gr.Interface(
23
+ fn=predict,
24
+ inputs=[
25
+ gr.Textbox(lines=5, label="Input Text"),
26
+ gr.Dropdown(models, label="Model")
27
+ ],
28
+ outputs='text',
29
  )
30
 
31
+ demo.launch()