main.py

from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from transformers import AutoTokenizer, AutoModelForTokenClassification, pipeline
import re
from typing import List, Dict, Optional, Any
import uvicorn
from hanspell import spell_checker

app = FastAPI(title="개인정보 필터링 & 오타 교정 API")

# 공개 한국어 NER 모델
model_name = "Leo97/KoELECTRA-small-v3-modu-ner"

print(f"모델 로딩 중: {model_name}")
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForTokenClassification.from_pretrained(model_name)
ner_pipeline = pipeline("ner", model=model, tokenizer=tokenizer, aggregation_strategy="simple")
print("✅ 모델 로드 완료!")

class TextRequest(BaseModel):
    text: str
    fix_spelling: Optional[bool] = True

class DetectedEntity(BaseModel):
    text: str
    type: str
    start: int
    end: int
    score: Optional[float] = None

class SpellingError(BaseModel):
    original: str
    corrected: str
    position: int

class FilteredResponse(BaseModel):
    original_text: str
    corrected_text: Optional[str] = None
    filtered_text: str
    detected_entities: List[DetectedEntity]
    spelling_errors: Optional[List[SpellingError]] = None

# 한글 숫자 매핑
KOREAN_NUMBER_MAP = {
    '공': '0', '영': '0',
    '일': '1', '하나': '1',
    '이': '2', '둘': '2',
    '삼': '3', '셋': '3',
    '사': '4', '넷': '4',
    '오': '5', '다섯': '5',
    '육': '6', '여섯': '6',
    '칠': '7', '일곱': '7',
    '팔': '8', '여덟': '8',
    '구': '9', '아홉': '9'
}

def korean_to_digit(text: str) -> str:
    """한글 숫자를 아라비아 숫자로 변환"""
    result = text
    for korean, digit in KOREAN_NUMBER_MAP.items():
        result = result.replace(korean, digit)
    return result

def extract_korean_phone_numbers(text: str) -> List[Dict[str, Any]]:
    """한글로 된 전화번호 패턴 추출 (공일공-구삼육팔-일팔삼이 등)"""
    entities = []
    
    # 한글 숫자 패턴 (공, 일, 이, 삼 등)
    korean_digits = r'[공영일이삼사오육칠팔구하나둘셋넷다섯여섯일곱여덟아홉]'
    
    # 전화번호 패턴: 한글숫자{2,3}-한글숫자{3,4}-한글숫자{4}
    patterns = [
        f'{korean_digits}{{2,3}}[-\s]{korean_digits}{{3,4}}[-\s]{korean_digits}{{4}}',
        f'{korean_digits}{{2,3}}{korean_digits}{{3,4}}{korean_digits}{{4}}',
        # 띄어쓰기로 구분된 경우
        f'{korean_digits}{{2,3}}\s{korean_digits}{{3,4}}\s{korean_digits}{{4}}',
    ]
    
    for pattern in patterns:
        matches = re.finditer(pattern, text)
        for match in matches:
            matched_text = match.group()
            # 한글을 숫자로 변환해서 검증
            converted = korean_to_digit(matched_text)
            # 숫자로 변환 후 전화번호 형태인지 확인
            digit_only = re.sub(r'[^0-9]', '', converted)
            if len(digit_only) >= 10 and len(digit_only) <= 11:
                entities.append({
                    "text": matched_text,
                    "type": "전화번호(한글)",
                    "start": match.start(),
                    "end": match.end(),
                    "converted": converted
                })
    
    return entities

def extract_mixed_phone_numbers(text: str) -> List[Dict[str, Any]]:
    """한글+숫자 혼합 전화번호 추출 (공1공-1234-5678 등)"""
    entities = []
    
    korean_digits = r'[공영일이삼사오육칠팔구하나둘셋넷다섯여섯일곱여덟아홉]'
    
    # 한글과 숫자가 섞인 패턴
    mixed_pattern = f'[{korean_digits}0-9]{{2,3}}[-\s]?[{korean_digits}0-9]{{3,4}}[-\s]?[{korean_digits}0-9]{{4}}'
    
    matches = re.finditer(mixed_pattern, text)
    for match in matches:
        matched_text = match.group()
        # 한글이 포함되어 있는지 확인
        if re.search(korean_digits, matched_text):
            converted = korean_to_digit(matched_text)
            digit_only = re.sub(r'[^0-9]', '', converted)
            if len(digit_only) >= 10 and len(digit_only) <= 11:
                entities.append({
                    "text": matched_text,
                    "type": "전화번호(혼합)",
                    "start": match.start(),
                    "end": match.end(),
                    "converted": converted
                })
    
    return entities

def check_spelling(text: str) -> Dict[str, Any]:
    """한국어 맞춤법 검사"""
    try:
        result = spell_checker.check(text)
        
        corrections = []
        corrected_text = result.checked
        
        if text != corrected_text:
            original_words = text.split()
            corrected_words = corrected_text.split()
            
            for i, (orig, corr) in enumerate(zip(original_words, corrected_words)):
                if orig != corr:
                    corrections.append({
                        "original": orig,
                        "corrected": corr,
                        "position": i
                    })
        
        return {
            "corrected_text": corrected_text,
            "errors": corrections
        }
    except Exception as e:
        print(f"맞춤법 검사 오류: {e}")
        return {
            "corrected_text": text,
            "errors": []
        }

def extract_phone_numbers(text: str) -> List[Dict[str, Any]]:
    """일반 숫자 전화번호 패턴 추출"""
    patterns = [
        r'\d{2,3}[-\s]?\d{3,4}[-\s]?\d{4}',
        r'\d{10,11}',
    ]
    
    entities = []
    for pattern in patterns:
        matches = re.finditer(pattern, text)
        for match in matches:
            entities.append({
                "text": match.group(),
                "type": "전화번호",
                "start": match.start(),
                "end": match.end()
            })
    return entities

def extract_email(text: str) -> List[Dict[str, Any]]:
    """이메일 패턴 추출"""
    pattern = r'\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Z|a-z]{2,}\b'
    entities = []
    matches = re.finditer(pattern, text)
    for match in matches:
        entities.append({
            "text": match.group(),
            "type": "이메일",
            "start": match.start(),
            "end": match.end()
        })
    return entities

def mask_text(text: str, mask_char: str = "0") -> str:
    """텍스트를 마스킹 문자로 변환"""
    return ''.join(mask_char if c not in [' ', '-', '@', '.'] else c for c in text)

def filter_sensitive_info(text: str, fix_spelling: bool = False,  mask_mapping: Dict[str, str] = None) -> Dict[str, Any]:
    """민감 정보 필터링 (000으로 마스킹) + 오타 교정"""
    
        # 기본 마스킹 매핑 설정
    if mask_mapping is None:
        mask_mapping = {
            "이름": "*",
            "지명": "#",
            "조직명": "@",
            "전화번호": "0",
            "이메일": "X",
            "default": "0"  # 기본값
        }
        
    # 1. 오타 교정 (옵션)
    spelling_result = None
    corrected_text = None
    if fix_spelling:
        spelling_result = check_spelling(text)
        corrected_text = spelling_result['corrected_text']
        text_for_ner = corrected_text
    else:
        text_for_ner = text
    
    all_entities = []
    
    # 2. NER 모델로 엔티티 추출
    try:
        ner_results = ner_pipeline(text_for_ner)
        
        entity_type_mapping = {
            "PER": "이름",
            "PERSON": "이름", 
            "PS": "이름",
            "LOC": "지명",
            "LOCATION": "지명",
            "LC": "지명",
            "ORG": "조직명",
            "ORGANIZATION": "조직명",
            "OG": "조직명",
        }
        
        for entity in ner_results:
            entity_type = entity.get('entity_group', '').upper()
            if not entity_type and 'entity' in entity:
                entity_type = entity['entity'].split('-')[-1].upper()
            
            if entity_type in entity_type_mapping:
                all_entities.append({
                    "text": str(entity['word']).strip(),
                    "type": entity_type_mapping[entity_type],
                    "start": int(entity['start']),
                    "end": int(entity['end']),
                    "score": float(round(entity.get('score', 0), 3))
                })
    except Exception as e:
        print(f"NER 처리 오류: {e}")
    
    # 3. 패턴 기반 추출
    phone_entities = extract_phone_numbers(text_for_ner)
    email_entities = extract_email(text_for_ner)
    korean_phone_entities = extract_korean_phone_numbers(text_for_ner)
    mixed_phone_entities = extract_mixed_phone_numbers(text_for_ner)
    
    all_entities.extend(phone_entities)
    all_entities.extend(email_entities)
    all_entities.extend(korean_phone_entities)
    all_entities.extend(mixed_phone_entities)
    
    # 4. 중복 제거 (겹치는 범위 제거)
    unique_entities = []
    seen_positions = set()
    
    # 시작 위치 기준으로 정렬
    sorted_entities = sorted(all_entities, key=lambda x: (x['start'], x['end']))
    
    for entity in sorted_entities:
        # 현재 엔티티가 기존 엔티티와 겹치는지 확인
        overlap = False
        for seen_start, seen_end in seen_positions:
            if not (entity['end'] <= seen_start or entity['start'] >= seen_end):
                overlap = True
                break
        
        if not overlap:
            seen_positions.add((entity['start'], entity['end']))
            unique_entities.append(entity)
    
    # 5. 뒤에서부터 정렬
    unique_entities_sorted = sorted(unique_entities, key=lambda x: x['start'], reverse=True)
    
    # 6. 텍스트 필터링
    filtered_text = text_for_ner
    for entity in unique_entities_sorted:
        start = entity['start']
        end = entity['end']
        original_text = filtered_text[start:end]
        
        # 000으로 마스킹
        mask_char = mask_mapping.get(entity['type'], mask_mapping['default'])
        masked = mask_text(original_text, mask_char)
        filtered_text = filtered_text[:start] + masked + filtered_text[end:]
    
    result = {
        "original_text": text,
        "filtered_text": filtered_text,
        "detected_entities": sorted(unique_entities, key=lambda x: x['start']),
        "mask_mapping_used": mask_mapping
    }
    
    # 오타 교정 결과 추가
    if fix_spelling and spelling_result:
        result["corrected_text"] = corrected_text
        result["spelling_errors"] = spelling_result['errors']
    
    return result

@app.post("/filter", response_model=FilteredResponse)
async def filter_text(request: TextRequest):
    """
    입력된 텍스트에서 개인정보를 000으로 마스킹합니다.
    
    - **text**: 필터링할 텍스트
    - **fix_spelling**: True로 설정하면 오타도 함께 교정합니다 (기본값: False)
    
    지원 패턴:
    - 일반 전화번호: 010-1234-5678
    - 한글 전화번호: 공일공-구삼육팔-일팔삼이
    - 혼합 전화번호: 공1공-1234-5678
    - 이메일: test@example.com
    - 이름, 지명, 조직명 (NER 모델)
    """
    try:
        if not request.text.strip():
            raise HTTPException(status_code=400, detail="텍스트가 비어있습니다.")
        
        result = filter_sensitive_info(request.text, request.fix_spelling)
        return result
    
    except Exception as e:
        print(f"오류 발생: {e}")
        raise HTTPException(status_code=500, detail=f"처리 중 오류 발생: {str(e)}")

@app.get("/")
async def root():
    return {
        "message": "개인정보 필터링 & 오타 교정 API",
        "model": model_name,
        "features": [
            "일반 전화번호 필터링 (010-1234-5678)",
            "한글 전화번호 필터링 (공일공-구삼육팔-일팔삼이)",
            "혼합 전화번호 필터링 (공1공-1234-5678)",
            "이메일 필터링",
            "이름, 지명, 조직명 필터링 (NER)",
            "한국어 맞춤법 검사"
        ],
        "endpoints": {
            "filter": "/filter (POST)",
            "docs": "/docs",
            "health": "/health"
        }
    }

@app.get("/health")
async def health_check():
    return {"status": "healthy", "model": model_name}

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=8000)