openfree commited on
Commit
d1c3219
·
verified ·
1 Parent(s): e4aee44

Update src/main.py

Browse files
Files changed (1) hide show
  1. src/main.py +4 -9
src/main.py CHANGED
@@ -2,7 +2,7 @@ import display_gloss as dg
2
  import synonyms_preprocess as sp
3
  from NLP_Spacy_base_translator import NlpSpacyBaseTranslator
4
  from flask import Flask, render_template, Response, request
5
- from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqTranslation
6
  import torch
7
  import os
8
 
@@ -22,8 +22,8 @@ os.environ['CUDA_VISIBLE_DEVICES'] = ''
22
 
23
  # Load pre-trained Korean-English translation model
24
  model_name = "Helsinki-NLP/opus-mt-ko-en"
25
- tokenizer = AutoTokenizer.from_pretrained(model_name, cache_dir=cache_dir)
26
- model = AutoModelForSeq2SeqTranslation.from_pretrained(model_name, cache_dir=cache_dir)
27
  model = model.to(device)
28
 
29
  nlp, dict_docs_spacy = sp.load_spacy_values()
@@ -31,12 +31,11 @@ dataset, list_2000_tokens = dg.load_data()
31
 
32
  def translate_korean_to_english(text):
33
  try:
34
- # Check if input is Korean
35
  if any('\u3131' <= char <= '\u318F' or '\uAC00' <= char <= '\uD7A3' for char in text):
36
  inputs = tokenizer(text, return_tensors="pt", padding=True)
37
  outputs = model.generate(**inputs)
38
  translation = tokenizer.decode(outputs[0], skip_special_tokens=True)
39
- print(f"Translated text: {translation}") # Debug log
40
  return translation
41
  return text
42
  except Exception as e:
@@ -52,18 +51,14 @@ def result():
52
  if request.method == 'POST':
53
  input_text = request.form['inputSentence']
54
  try:
55
- # Translate to English
56
  english_text = translate_korean_to_english(input_text)
57
 
58
- # Check if translation failed
59
  if english_text == input_text and any('\u3131' <= char <= '\u318F' or '\uAC00' <= char <= '\uD7A3' for char in input_text):
60
  raise Exception("Translation failed")
61
 
62
- # Convert to ASL gloss
63
  eng_to_asl_translator = NlpSpacyBaseTranslator(sentence=english_text)
64
  generated_gloss = eng_to_asl_translator.translate_to_gloss()
65
 
66
- # Process gloss
67
  gloss_list_lower = [gloss.lower() for gloss in generated_gloss.split() if gloss.isalnum()]
68
  gloss_sentence_before_synonym = " ".join(gloss_list_lower)
69
 
 
2
  import synonyms_preprocess as sp
3
  from NLP_Spacy_base_translator import NlpSpacyBaseTranslator
4
  from flask import Flask, render_template, Response, request
5
+ from transformers import MarianMTModel, MarianTokenizer
6
  import torch
7
  import os
8
 
 
22
 
23
  # Load pre-trained Korean-English translation model
24
  model_name = "Helsinki-NLP/opus-mt-ko-en"
25
+ tokenizer = MarianTokenizer.from_pretrained(model_name, cache_dir=cache_dir)
26
+ model = MarianMTModel.from_pretrained(model_name, cache_dir=cache_dir)
27
  model = model.to(device)
28
 
29
  nlp, dict_docs_spacy = sp.load_spacy_values()
 
31
 
32
  def translate_korean_to_english(text):
33
  try:
 
34
  if any('\u3131' <= char <= '\u318F' or '\uAC00' <= char <= '\uD7A3' for char in text):
35
  inputs = tokenizer(text, return_tensors="pt", padding=True)
36
  outputs = model.generate(**inputs)
37
  translation = tokenizer.decode(outputs[0], skip_special_tokens=True)
38
+ print(f"Translated text: {translation}")
39
  return translation
40
  return text
41
  except Exception as e:
 
51
  if request.method == 'POST':
52
  input_text = request.form['inputSentence']
53
  try:
 
54
  english_text = translate_korean_to_english(input_text)
55
 
 
56
  if english_text == input_text and any('\u3131' <= char <= '\u318F' or '\uAC00' <= char <= '\uD7A3' for char in input_text):
57
  raise Exception("Translation failed")
58
 
 
59
  eng_to_asl_translator = NlpSpacyBaseTranslator(sentence=english_text)
60
  generated_gloss = eng_to_asl_translator.translate_to_gloss()
61
 
 
62
  gloss_list_lower = [gloss.lower() for gloss in generated_gloss.split() if gloss.isalnum()]
63
  gloss_sentence_before_synonym = " ".join(gloss_list_lower)
64