import gradio as gr
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
from peft import PeftModel
import spaces

# 1. 모델 ID 설정
base_id = "Qwen/Qwen2.5-7B-Instruct"
adapter_id = "epinfomax/BizFlow-Summarizer-Ko"

# 2. 모델 로드
print(f"🚀 모델 로딩 중...")
tokenizer = AutoTokenizer.from_pretrained(base_id)
model = AutoModelForCausalLM.from_pretrained(base_id, torch_dtype=torch.float16)
model = PeftModel.from_pretrained(model, adapter_id)

# 3. 추론 함수
@spaces.GPU
def summarize(text):
    device = "cuda"
    model.to(device)
    
    messages = [
        {"role": "system", "content": "당신은 비즈니스 문서를 전문적으로 요약하는 AI 어시스턴트입니다. 핵심 내용을 명확하게 요약해 주세요."},
        {"role": "user", "content": text}
    ]
    
    input_text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    inputs = tokenizer([input_text], return_tensors="pt").to(device)
    
    with torch.no_grad():
        outputs = model.generate(
            **inputs, 
            max_new_tokens=512,
            temperature=0.3,
            repetition_penalty=1.1
        )
    
    generated_tokens = outputs[:, inputs.input_ids.shape[1]:]
    result = tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)
    
    return result[0]

# ★ 예시용 긴 문장 데이터 (실제 기사 스타일)
example_text = """삼성전자가 2024년 4분기 연결 기준 매출 67조 7800억 원, 영업이익 2조 8200억 원을 기록했다고 31일 확정 공시했다.

전년 동기 대비 매출은 3.81% 감소했고, 영업이익은 34.4% 줄었다. 메모리 반도체 시황 회복으로 전분기 대비 매출과 영업이익은 개선됐다. 
부문별로 보면 반도체(DS) 부문은 매출 21조 7000억 원, 영업손실 2조 1800억 원을 기록했다. 
D램은 재고 수준이 큰 폭으로 개선된 가운데 흑자 전환에 성공했다. 파운드리는 시장 수요 침체로 실적 부진이 지속됐다.

모바일(MX)과 네트워크 부문은 매출 25조 400억 원, 영업이익 2조 7300억 원을 달성했다. 
신제품 출시 효과가 둔화되며 스마트폰 판매량은 줄었지만, 태블릿과 웨어러블 제품 판매가 견조했다. 
영상디스플레이(VD)와 가전 부문은 매출 14조 2200억 원, 영업손실 500억 원을 기록했다. 
글로벌 TV 수요 정체에도 프리미엄 제품 중심으로 판매 구조를 개선했다.

삼성전자는 "올해 1분기에는 IT 시황 회복세가 기대되는 만큼, 고부가가치 제품 판매를 확대해 수익성 개선에 주력할 계획"이라고 밝혔다. 
특히 생성형 AI용 고대역폭메모리(HBM) 등 선단 공정 제품 판매를 대폭 늘릴 방침이다."""

# 4. 웹 인터페이스 정의
iface = gr.Interface(
    fn=summarize,
    inputs=gr.Textbox(
        lines=15, 
        placeholder="요약할 문서를 여기에 붙여넣으세요...", 
        label="입력 문서"
    ),
    outputs=gr.Textbox(
        lines=10,             # ★ 결과창 높이를 10줄로 고정하여 시원하게 만듦
        label="요약 결과",
        #show_copy_button=True # ★ 복사 버튼 추가 (편의성)
    ),
    title="BizFlow 문서 요약기",
    description="Qwen2.5-7B + 파인튜닝(LoRA) 모델 테스트 데모입니다. (ZeroGPU)",
    examples=[[example_text]] # ★ 위에서 정의한 긴 문장을 예시로 넣음
)

if __name__ == "__main__":
    iface.launch()