richylyq commited on
Commit
bccf7be
·
1 Parent(s): 5b11a3e
Files changed (1) hide show
  1. app.py +16 -6
app.py CHANGED
@@ -126,22 +126,32 @@ def opus_trans(article, target_language):
126
 
127
 
128
  def nllb_trans(article, target_language):
 
 
129
  tokenizer = AutoTokenizer.from_pretrained("facebook/nllb-200-distilled-600M")
130
  model = AutoModelForSeq2SeqLM.from_pretrained("facebook/nllb-200-distilled-600M")
131
  inputs = tokenizer(article, return_tensors="pt")
132
 
133
  if target_language == "English":
134
  target_lang = "eng_Latn"
 
135
  elif target_language == "Chinese":
136
  target_lang = "zho_Hans"
 
137
 
138
- translated_tokens = model.generate(
139
- **inputs,
140
- forced_bos_token_id=tokenizer.lang_code_to_id[target_lang],
141
- max_length=30,
142
- )
 
 
 
 
 
 
143
 
144
- return tokenizer.batch_decode(translated_tokens, skip_special_tokens=True)[0]
145
 
146
 
147
  def translate(article, toolkit, target_language):
 
126
 
127
 
128
  def nllb_trans(article, target_language):
129
+ result_lang = detect_lang(article)
130
+
131
  tokenizer = AutoTokenizer.from_pretrained("facebook/nllb-200-distilled-600M")
132
  model = AutoModelForSeq2SeqLM.from_pretrained("facebook/nllb-200-distilled-600M")
133
  inputs = tokenizer(article, return_tensors="pt")
134
 
135
  if target_language == "English":
136
  target_lang = "eng_Latn"
137
+ target_language = "en"
138
  elif target_language == "Chinese":
139
  target_lang = "zho_Hans"
140
+ target_language = "zh"
141
 
142
+ if result_lang != target_language:
143
+ translated_tokens = model.generate(
144
+ **inputs,
145
+ forced_bos_token_id=tokenizer.lang_code_to_id[target_lang],
146
+ max_length=30,
147
+ )
148
+ translated = tokenizer.batch_decode(
149
+ translated_tokens, skip_special_tokens=True
150
+ )[0]
151
+ else:
152
+ translated = "Error: You chose the same language as the article detected language. Please reselect language and try again."
153
 
154
+ return translated
155
 
156
 
157
  def translate(article, toolkit, target_language):