from transformers import PreTrainedTokenizerFast, BartForConditionalGeneration
import torch

"""
[설계 의도]
텍스트 요약은 생성 길이와 하이퍼파라미터를 세밀하게 제어해야 하므로, Pipeline 대신 AutoClass API를 사용했습니다.
한국어 처리 정확도를 높이기 위해 KoBART의 BPE 기반 토크나이저를 직접 적용해 형태소 손실을 최소화했습니다.
"""

model_name = "ainize/kobart-news"
tokenizer = PreTrainedTokenizerFast.from_pretrained(model_name)
model = BartForConditionalGeneration.from_pretrained(model_name)

if torch.cuda.is_available():
    model = model.to('cuda')

def summarize_review(text, min_len, max_len):
    # 텐서 변환
    input_ids = tokenizer.encode(text, return_tensors="pt")
    if torch.cuda.is_available():
        input_ids = input_ids.to('cuda')
        
    # 명시적 생성 제어
    summary_ids = model.generate(
        input_ids=input_ids,
        bos_token_id=model.config.bos_token_id,
        eos_token_id=model.config.eos_token_id,
        length_penalty=1.0,       # 2.0에서 1.0으로 낮춰서 불필요하게 길게 쓰는 것을 방지
        max_length=int(max_len),
        min_length=int(min_len),
        num_beams=4,
        early_stopping=True       # 문장이 마침표로 완성되면 제한 길이에 도달 안 해도 즉시 종료!
    )
    
    # 자연어로 디코딩
    answer = tokenizer.decode(summary_ids[0], skip_special_tokens=True)
    
    # 문장이 중간에 잘린 경우를 대비한 깔끔한 후처리
    if not answer.endswith(('.', '!', '?')) and '.' in answer:
        answer = answer[:answer.rfind('.') + 1]
        
    # 자주 틀리는 오타나 어색한 어미를 파이썬 문자열 교체로 강제 교정합니다.
    answer = answer.replace("네다.", "네요.")
    answer = answer.replace("너무하네다", "너무하네요")
    answer = answer.replace("뿐이.", "뿐입니다.")
    answer = answer.replace("로썸틱", "로맨틱") # 만약 자주 나오는 이상한 단어가 있다면 이렇게 추가
    
    return answer