Update src/main.py

src/main.py

@@ -2,33 +2,38 @@ import display_gloss as dg
 import synonyms_preprocess as sp
 from NLP_Spacy_base_translator import NlpSpacyBaseTranslator 
 from flask import Flask, render_template, Response, request
-from transformers import pipeline
+from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqGeneration
 import torch
 import os
 
 app = Flask(__name__)
 app.config['TITLE'] = 'ASL Translator'
 
-# Set cache directory
-os.environ['TRANSFORMERS_CACHE'] = '/tmp/transformers_cache'
-os.makedirs('/tmp/transformers_cache', exist_ok=True)
+# 캐시 디렉토리 설정
+cache_dir = "/tmp/huggingface"
+if not os.path.exists(cache_dir):
+    os.makedirs(cache_dir, exist_ok=True)
+os.environ['TRANSFORMERS_CACHE'] = cache_dir
+os.environ['HF_HOME'] = cache_dir
 
-# Force CPU usage
+# CPU 설정
 device = torch.device('cpu')
 os.environ['CUDA_VISIBLE_DEVICES'] = ''
 
-# Initialize translator with local cache
-translator = pipeline("translation", 
-                     model="Helsinki-NLP/opus-mt-ko-en", 
-                     device=device,
-                     model_kwargs={"cache_dir": "/tmp/transformers_cache"})
+# 번역 모델 초기화
+model_name = "Helsinki-NLP/opus-mt-ko-en"
+tokenizer = AutoTokenizer.from_pretrained(model_name, cache_dir=cache_dir)
+model = AutoModelForSeq2SeqGeneration.from_pretrained(model_name, cache_dir=cache_dir)
+model = model.to(device)
 
 nlp, dict_docs_spacy = sp.load_spacy_values()
 dataset, list_2000_tokens = dg.load_data()
 
 def translate_korean_to_english(text):
     if any('\u3131' <= char <= '\u318F' or '\uAC00' <= char <= '\uD7A3' for char in text):
-        translation = translator(text)[0]['translation_text']
+        inputs = tokenizer(text, return_tensors="pt", padding=True)
+        outputs = model.generate(**inputs)
+        translation = tokenizer.decode(outputs[0], skip_special_tokens=True)
         return translation
     return text