LAP-DEV commited on
Commit
40135c8
·
verified ·
1 Parent(s): ff49e9c

Update modules/translation/translation_base.py

Browse files
modules/translation/translation_base.py CHANGED
@@ -136,7 +136,8 @@ class TranslationBase(ABC):
136
  finally:
137
  self.release_cuda_memory()
138
 
139
- def translate_text(input_list_dict: list,
 
140
  model_size: str,
141
  src_lang: str,
142
  tgt_lang: str,
@@ -170,13 +171,13 @@ class TranslationBase(ABC):
170
  List[dict] with translation
171
  """
172
  try:
173
- cache_parameters(model_size=model_size,src_lang=src_lang,tgt_lang=tgt_lang,max_length=max_length,add_timestamp=add_timestamp)
174
- update_model(model_size=model_size,src_lang=src_lang,tgt_lang=tgt_lang,progress=progress)
175
 
176
  total_progress = len(input_list_dict)
177
  for index, dic in enumerate(input_list_dict):
178
  progress(index / total_progress, desc="Translating..")
179
- translated_text = translate(dic["text"], max_length=max_length)
180
  dic["text"] = translated_text
181
 
182
  return input_list_dict
@@ -184,7 +185,7 @@ class TranslationBase(ABC):
184
  except Exception as e:
185
  print(f"Error: {str(e)}")
186
  finally:
187
- release_cuda_memory()
188
 
189
  @staticmethod
190
  def get_device():
 
136
  finally:
137
  self.release_cuda_memory()
138
 
139
+ def translate_text(self,
140
+ input_list_dict: list,
141
  model_size: str,
142
  src_lang: str,
143
  tgt_lang: str,
 
171
  List[dict] with translation
172
  """
173
  try:
174
+ self.cache_parameters(model_size=model_size,src_lang=src_lang,tgt_lang=tgt_lang,max_length=max_length,add_timestamp=add_timestamp)
175
+ self.update_model(model_size=model_size,src_lang=src_lang,tgt_lang=tgt_lang,progress=progress)
176
 
177
  total_progress = len(input_list_dict)
178
  for index, dic in enumerate(input_list_dict):
179
  progress(index / total_progress, desc="Translating..")
180
+ translated_text = self.translate(dic["text"], max_length=max_length)
181
  dic["text"] = translated_text
182
 
183
  return input_list_dict
 
185
  except Exception as e:
186
  print(f"Error: {str(e)}")
187
  finally:
188
+ self.release_cuda_memory()
189
 
190
  @staticmethod
191
  def get_device():