from flask import Flask, request, jsonify
from transformers import pipeline

# Flask 애플리케이션 초기화
app = Flask(__name__)

# 문법 체크를 위한 'fill-mask' 파이프라인 로드
grammar_checker = pipeline('fill-mask', model='bert-base-uncased')

# 문법 체크 함수
def check_grammar(sentence):
    words = sentence.split()
    checked_sentence = []

    for i in range(len(words)):
        original_word = words[i]
        # 마스크된 문장 생성
        masked_sentence = ' '.join(words[:i] + ['[MASK]'] + words[i+1:])
        # 모델을 이용해 마스크된 단어 예측
        result = grammar_checker(masked_sentence)
        # 원래 단어와 가장 비슷한 예측 결과 선택
        best_prediction = min(result, key=lambda x: abs(len(x['token_str']) - len(original_word)))
        suggested_word = best_prediction['token_str']
        checked_sentence.append(suggested_word)

    return ' '.join(checked_sentence)

# 엔드포인트 정의
@app.route('/check_grammar', methods=['POST'])
def grammar_check_endpoint():
    data = request.json
    sentence = data.get('sentence', '')
    corrected_sentence = check_grammar(sentence)
    return jsonify({
        'original': sentence,
        'corrected': corrected_sentence
    })

# 애플리케이션 실행
if __name__ == '__main__':
    app.run(debug=True)