Spaces:
Building
Building
Update src/main.py
Browse files- src/main.py +2 -2
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,
|
6 |
import torch
|
7 |
import os
|
8 |
|
@@ -23,7 +23,7 @@ os.environ['CUDA_VISIBLE_DEVICES'] = ''
|
|
23 |
# 번역 모델 초기화
|
24 |
model_name = "Helsinki-NLP/opus-mt-ko-en"
|
25 |
tokenizer = AutoTokenizer.from_pretrained(model_name, cache_dir=cache_dir)
|
26 |
-
model =
|
27 |
model = model.to(device)
|
28 |
|
29 |
nlp, dict_docs_spacy = sp.load_spacy_values()
|
|
|
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 |
|
|
|
23 |
# 번역 모델 초기화
|
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()
|