openfree commited on
Commit
6ba0d40
·
verified ·
1 Parent(s): f687683

Update src/main.py

Browse files
Files changed (1) hide show
  1. 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, AutoModelForSeq2SeqGeneration
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 = AutoModelForSeq2SeqGeneration.from_pretrained(model_name, cache_dir=cache_dir)
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()