import re

def normalize_mongolian_text(text):
    # HTML 태그 제거
    text = re.sub(r'<[^>]+>', '', text)
    
    # 연속된 공백을 하나로
    text = re.sub(r'\s+', ' ', text)
    
    # 앞뒤 공백 제거
    text = text.strip()
    
    return text

# 사용 예시
raw_text = "  Сайн   байна уу?  <br>  "
clean_text = normalize_mongolian_text(raw_text)
print(clean_text)  # "Сайн байна уу?"

from transformers import AutoTokenizer, AutoModelForCausalLM

# XLM-RoBERTa 모델 로드
model_name = "xlm-roberta-base"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)

# 몽골어 텍스트 인코딩
mongolian_text = "Сайн байна уу?"
inputs = tokenizer(mongolian_text, return_tensors="pt")

# 모델 추론
outputs = model(**inputs)

from torch.utils.data import Dataset, DataLoader

class MongolianChatDataset(Dataset):
    def __init__(self, texts, tokenizer, max_length=128):
        self.texts = texts
        self.tokenizer = tokenizer
        self.max_length = max_length
    
    def __len__(self):
        return len(self.texts)
    
    def __getitem__(self, idx):
        text = self.texts[idx]
        encoding = self.tokenizer(
            text,
            max_length=self.max_length,
            padding='max_length',
            truncation=True,
            return_tensors='pt'
        )
        return {
            'input_ids': encoding['input_ids'].flatten(),
            'attention_mask': encoding['attention_mask'].flatten()
        }

# 데이터 로더 생성
train_dataset = MongolianChatDataset(train_texts, tokenizer)
train_loader = DataLoader(train_dataset, batch_size=16, shuffle=True)

from transformers import AdamW, get_linear_schedule_with_warmup

# 옵티마이저 설정
optimizer = AdamW(model.parameters(), lr=2e-5)

# 학습률 스케줄러
total_steps = len(train_loader) * num_epochs
warmup_steps = total_steps // 10

scheduler = get_linear_schedule_with_warmup(
    optimizer,
    num_warmup_steps=warmup_steps,
    num_training_steps=total_steps
)

import torch
from tqdm import tqdm

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model.to(device)

for epoch in range(num_epochs):
    model.train()
    total_loss = 0
    
    progress_bar = tqdm(train_loader, desc=f'Epoch {epoch+1}')
    
    for batch in progress_bar:
        # 데이터를 GPU로 이동
        input_ids = batch['input_ids'].to(device)
        attention_mask = batch['attention_mask'].to(device)
        
        # Forward pass
        outputs = model(
            input_ids=input_ids,
            attention_mask=attention_mask,
            labels=input_ids
        )
        
        loss = outputs.loss
        total_loss += loss.item()
        
        # Backward pass
        optimizer.zero_grad()
        loss.backward()
        
        # 그래디언트 클리핑 (안정적인 학습을 위해)
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        
        optimizer.step()
        scheduler.step()
        
        progress_bar.set_postfix({'loss': loss.item()})
    
    avg_loss = total_loss / len(train_loader)
    print(f'Epoch {epoch+1} - Average Loss: {avg_loss:.4f}')

# TensorBoard 사용 예시
from torch.utils.tensorboard import SummaryWriter

writer = SummaryWriter('runs/mongolian_chatbot')

# 훈련 루프 안에서
writer.add_scalar('Loss/train', loss.item(), global_step)
writer.add_scalar('Learning_Rate', scheduler.get_last_lr()[0], global_step)

# 훈련 종료 후
writer.close()

# 역번역을 이용한 데이터 증강 예시
from transformers import MarianMTModel, MarianTokenizer

def back_translate(text, src_lang='mn', pivot_lang='en'):
    # 몽골어 → 영어
    model_name_mn_en = f'Helsinki-NLP/opus-mt-{src_lang}-{pivot_lang}'
    tokenizer_mn_en = MarianTokenizer.from_pretrained(model_name_mn_en)
    model_mn_en = MarianMTModel.from_pretrained(model_name_mn_en)
    
    translated = model_mn_en.generate(
        **tokenizer_mn_en(text, return_tensors="pt", padding=True)
    )
    english_text = tokenizer_mn_en.decode(translated[0], skip_special_tokens=True)
    
    # 영어 → 몽골어
    model_name_en_mn = f'Helsinki-NLP/opus-mt-{pivot_lang}-{src_lang}'
    tokenizer_en_mn = MarianTokenizer.from_pretrained(model_name_en_mn)
    model_en_mn = MarianMTModel.from_pretrained(model_name_en_mn)
    
    back_translated = model_en_mn.generate(
        **tokenizer_en_mn(english_text, return_tensors="pt", padding=True)
    )
    
    return tokenizer_en_mn.decode(back_translated[0], skip_special_tokens=True)

# 사용 예시
original = "Сайн байна уу?"
augmented = back_translate(original)
print(f"원본: {original}")
print(f"증강: {augmented}")

# Few-Shot 프롬프트 예시
prompt = """
다음은 몽골어 대화 예시입니다:

사용자: Сайн байна уу?
챗봇: Сайн байна! Танд хэрхэн туслах вэ?

사용자: Би монгол хэл сурч байна.
챗봇: Гайхалтай! Монгол хэл сурах нь маш сонирхолтой.

사용자: Өнөөдөр цаг агаар ямар байна?
챗봇:
"""

# 모델이 이 패턴을 학습해서 적절한 답변을 생성

# BLEU Score 계산 예시
from nltk.translate.bleu_score import sentence_bleu

reference = [['Сайн', 'байна', 'уу']]
candidate = ['Сайн', 'байна', 'уу']

score = sentence_bleu(reference, candidate)
print(f'BLEU Score: {score:.4f}')

# Perplexity 계산
import torch
import math

def calculate_perplexity(model, test_loader):
    model.eval()
    total_loss = 0
    total_tokens = 0
    
    with torch.no_grad():
        for batch in test_loader:
            outputs = model(**batch)
            loss = outputs.loss
            total_loss += loss.item() * batch['input_ids'].size(0)
            total_tokens += batch['input_ids'].size(0)
    
    avg_loss = total_loss / total_tokens
    perplexity = math.exp(avg_loss)
    
    return perplexity

# FastAPI를 이용한 간단한 챗봇 API
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

app = FastAPI()

# 모델 로드 (서버 시작 시 한 번만)
model_name = "your-mongolian-chatbot-model"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model.to(device)

class ChatRequest(BaseModel):
    message: str
    user_id: str = None
    context: list = []

class ChatResponse(BaseModel):
    response: str
    confidence: float

@app.post("/chat", response_model=ChatResponse)
async def chat(request: ChatRequest):
    try:
        # 입력 텍스트 인코딩
        inputs = tokenizer(
            request.message,
            return_tensors="pt",
            padding=True,
            truncation=True
        ).to(device)
        
        # 응답 생성
        with torch.no_grad():
            outputs = model.generate(
                **inputs,
                max_length=100,
                num_return_sequences=1,
                temperature=0.7,
                top_p=0.9,
                do_sample=True
            )
        
        # 디코딩
        response_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
        
        return ChatResponse(
            response=response_text,
            confidence=0.85  # 실제로는 모델의 확신도 계산
        )
    
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

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

# 모델 양자화 예시
import torch
from transformers import AutoModelForCausalLM

# 원본 모델 로드
model = AutoModelForCausalLM.from_pretrained("mongolian-chatbot")

# 동적 양자화 적용
quantized_model = torch.quantization.quantize_dynamic(
    model,
    {torch.nn.Linear},  # 양자화할 레이어 타입
    dtype=torch.qint8   # INT8로 변환
)

# 모델 크기 비교
original_size = sum(p.numel() for p in model.parameters()) * 4 / (1024**2)
quantized_size = sum(p.numel() for p in quantized_model.parameters()) / (1024**2)

print(f"원본 모델: {original_size:.2f} MB")
print(f"양자화 모델: {quantized_size:.2f} MB")
print(f"크기 감소: {(1 - quantized_size/original_size)*100:.1f}%")