davanstrien HF staff commited on
Commit
3f23d73
1 Parent(s): 6bcde50

Refactor code to improve performance and readability

Browse files
Files changed (1) hide show
  1. app.py +13 -10
app.py CHANGED
@@ -14,11 +14,17 @@ flores_codes = list(code_mapping.keys())
14
 
15
  def load_model():
16
  model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_NAME).to(device)
17
- tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
18
- return model, tokenizer
19
 
20
 
21
- model, tokenizer = load_model()
 
 
 
 
 
 
 
22
 
23
 
24
  @spaces.GPU
@@ -29,12 +35,10 @@ def translate(
29
  window_size: int = 800,
30
  overlap_size: int = 200,
31
  ):
 
 
32
  input_tokens = (
33
- tokenizer(text, return_tensors="pt", src_lang=code_mapping[src_lang])
34
- .input_ids[0]
35
- .cpu()
36
- .numpy()
37
- .tolist()
38
  )
39
  translated_chunks = []
40
 
@@ -51,8 +55,7 @@ def translate(
51
  )
52
  translated_chunks.append(translated_chunk)
53
 
54
- translated_text = " ".join(translated_chunks)
55
- return translated_text
56
 
57
 
58
  description = """
 
14
 
15
  def load_model():
16
  model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_NAME).to(device)
17
+ return model
 
18
 
19
 
20
+ model = load_model()
21
+
22
+
23
+ def load_tokenizer(src_lang, tgt_lang):
24
+ tokenizer = AutoTokenizer.from_pretrained(
25
+ MODEL_NAME, src_lang=code_mapping[src_lang], tgt_lang=code_mapping[tgt_lang]
26
+ )
27
+ return tokenizer
28
 
29
 
30
  @spaces.GPU
 
35
  window_size: int = 800,
36
  overlap_size: int = 200,
37
  ):
38
+ tokenizer = load_tokenizer(src_lang, tgt_lang)
39
+
40
  input_tokens = (
41
+ tokenizer(text, return_tensors="pt").input_ids[0].cpu().numpy().tolist()
 
 
 
 
42
  )
43
  translated_chunks = []
44
 
 
55
  )
56
  translated_chunks.append(translated_chunk)
57
 
58
+ return " ".join(translated_chunks)
 
59
 
60
 
61
  description = """