sudoping01 commited on
Commit
8948ad3
·
verified ·
1 Parent(s): 9331fff

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +6 -3
app.py CHANGED
@@ -1,4 +1,5 @@
1
  import gradio as gr
 
2
  from transformers import pipeline
3
  from typing import Union, List
4
 
@@ -8,8 +9,8 @@ class BambaraTranslator:
8
  "translation",
9
  model=model_name,
10
  max_length=512,
11
- truncation=True ,
12
- device=-1 # Force CPU
13
  )
14
  self.flores_codes = {
15
  "French": "fra_Latn",
@@ -17,6 +18,7 @@ class BambaraTranslator:
17
  "Bambara": "bam_Latn"
18
  }
19
 
 
20
  def translate(self, text: Union[str, List[str]], src_lang: str, tgt_lang: str) -> Union[str, List[str]]:
21
  source_lang = self.flores_codes[src_lang]
22
  target_lang = self.flores_codes[tgt_lang]
@@ -36,6 +38,7 @@ examples = [
36
  "The sun rises every morning to bring light to the world."
37
  ]
38
 
 
39
  def translate_text(src_lang, tgt_lang, text):
40
  if not text.strip():
41
  return "Please enter text to translate."
@@ -75,7 +78,7 @@ with gr.Blocks(title="Bambara Translator") as demo:
75
 
76
  gr.Markdown(
77
  """
78
- By [sudoping01](https://huggingface.co/sudoping01), from [sudoping01/nllb-bambara-v2](https://huggingface.co/sudoping01/nllb-bambara-v2). Fine-tuned on Metas NLLB, CC BY-NC 4.0, non-commercial.
79
  """
80
  )
81
 
 
1
  import gradio as gr
2
+ import spaces
3
  from transformers import pipeline
4
  from typing import Union, List
5
 
 
9
  "translation",
10
  model=model_name,
11
  max_length=512,
12
+ truncation=True,
13
+ device=0 if spaces.config.Hardware.current == "cuda" else -1
14
  )
15
  self.flores_codes = {
16
  "French": "fra_Latn",
 
18
  "Bambara": "bam_Latn"
19
  }
20
 
21
+ @spaces.GPU
22
  def translate(self, text: Union[str, List[str]], src_lang: str, tgt_lang: str) -> Union[str, List[str]]:
23
  source_lang = self.flores_codes[src_lang]
24
  target_lang = self.flores_codes[tgt_lang]
 
38
  "The sun rises every morning to bring light to the world."
39
  ]
40
 
41
+ @spaces.GPU
42
  def translate_text(src_lang, tgt_lang, text):
43
  if not text.strip():
44
  return "Please enter text to translate."
 
78
 
79
  gr.Markdown(
80
  """
81
+ By [sudoping01](https://huggingface.co/sudoping01), from [sudoping01/nllb-bambara-v2](https://huggingface.co/sudoping01/nllb-bambara-v2). Fine-tuned on Meta's NLLB, CC BY-NC 4.0, non-commercial.
82
  """
83
  )
84