| """한전 공문 검토·작성 API (FastAPI + onnxruntime) |
| |
| 환경 변수: |
| MODEL_REPO — HF 모델 repo (기본: onnx-community/EXAONE-3.5-2.4B-Instruct) |
| """ |
| import json |
| import os |
| import threading |
| import time |
| from pathlib import Path |
| from typing import Optional |
|
|
| import re |
| import tempfile |
| import xml.etree.ElementTree as ET |
| import zipfile |
|
|
| from fastapi import FastAPI, File, HTTPException, UploadFile |
| from fastapi.middleware.cors import CORSMiddleware |
| from fastapi.responses import FileResponse, StreamingResponse |
| from fastapi.staticfiles import StaticFiles |
| from pydantic import BaseModel |
|
|
| MODEL_REPO = os.getenv("MODEL_REPO", "onnx-community/EXAONE-3.5-2.4B-Instruct") |
| DATAS_DIR = Path(__file__).parent / "datas" |
| DOC_TYPES = {"공문-외부발송", "공문-내부", "보고서"} |
|
|
| np = None |
| ort = None |
|
|
| NUM_LAYERS = 30 |
| NUM_KV_HEADS = 8 |
| HEAD_DIM = 80 |
|
|
| app = FastAPI(title="공문 검토·작성 API", version="0.2") |
| app.add_middleware( |
| CORSMiddleware, |
| allow_origins=["*"], |
| allow_methods=["*"], |
| allow_headers=["*"], |
| ) |
|
|
|
|
| class Generator: |
| def __init__(self): |
| self.tokenizer = None |
| self.session = None |
| self.output_index = {} |
| self.model_dir = None |
| self.ready = False |
| self.status = "initializing" |
|
|
| def load(self): |
| import numpy as _np |
| import onnxruntime as _ort |
| from huggingface_hub import snapshot_download |
| from transformers import AutoTokenizer |
| global np, ort |
| np = _np |
| ort = _ort |
|
|
| try: |
| self.status = "downloading model" |
| self.model_dir = Path(snapshot_download( |
| MODEL_REPO, |
| allow_patterns=[ |
| "config.json", |
| "generation_config.json", |
| "tokenizer.json", |
| "tokenizer_config.json", |
| "special_tokens_map.json", |
| "onnx/model_q4.onnx", |
| "onnx/model_q4.onnx_data", |
| ], |
| )) |
| self.status = "loading tokenizer" |
| self.tokenizer = AutoTokenizer.from_pretrained( |
| str(self.model_dir), trust_remote_code=True |
| ) |
|
|
| self.status = "creating ONNX session" |
| so = ort.SessionOptions() |
| so.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL |
| try: |
| actual_cpus = len(os.sched_getaffinity(0)) |
| except AttributeError: |
| actual_cpus = os.cpu_count() or 2 |
| n_threads = int(os.getenv("OMP_NUM_THREADS", str(actual_cpus))) |
| so.intra_op_num_threads = max(1, n_threads) |
| so.inter_op_num_threads = 1 |
| print(f"[ONNX] threads: intra={so.intra_op_num_threads}, inter=1 (detected {actual_cpus} CPUs)") |
| self.session = ort.InferenceSession( |
| str(self.model_dir / "onnx" / "model_q4.onnx"), |
| sess_options=so, |
| providers=["CPUExecutionProvider"], |
| ) |
| self.output_index = { |
| o.name: i for i, o in enumerate(self.session.get_outputs()) |
| } |
| self.ready = True |
| self.status = "ready" |
| except Exception as e: |
| self.status = f"failed: {e}" |
| raise |
|
|
| def _sample(self, logits, temperature, top_p, top_k): |
| logits = logits.astype(np.float32) |
| if temperature <= 0: |
| return int(np.argmax(logits)) |
| logits = logits / temperature |
| if 0 < top_k < logits.size: |
| kth = np.partition(logits, -top_k)[-top_k] |
| logits = np.where(logits < kth, -np.inf, logits) |
| probs = np.exp(logits - logits.max()) |
| probs = probs / probs.sum() |
| if top_p < 1.0: |
| order = np.argsort(-probs) |
| sorted_p = probs[order] |
| cum = np.cumsum(sorted_p) |
| cutoff = np.searchsorted(cum, top_p) + 1 |
| keep_idx = order[:cutoff] |
| new_probs = np.zeros_like(probs) |
| new_probs[keep_idx] = probs[keep_idx] |
| probs = new_probs / new_probs.sum() |
| return int(np.random.choice(len(probs), p=probs)) |
|
|
| def _format_exaone(self, messages): |
| parts = [] |
| for m in messages: |
| role = m["role"] |
| content = m["content"] |
| if role == "user": |
| parts.append(f"[|{role}|]{content}\n") |
| else: |
| parts.append(f"[|{role}|]{content}[|endofturn|]\n") |
| parts.append("[|assistant|]") |
| return "".join(parts) |
|
|
| def stream(self, messages, max_new_tokens, temperature, top_p, top_k): |
| tok = self.tokenizer |
| try: |
| ids = tok.apply_chat_template( |
| messages, return_tensors="np", add_generation_prompt=True |
| ).astype(np.int64) |
| except Exception as e: |
| print(f"[WARN] apply_chat_template failed ({e}), using manual format") |
| prompt_str = self._format_exaone(messages) |
| ids = tok(prompt_str, return_tensors="np").input_ids.astype(np.int64) |
| input_ids = ids |
| seq_len = input_ids.shape[1] |
| attention_mask = np.ones((1, seq_len), dtype=np.int64) |
| position_ids = np.arange(seq_len, dtype=np.int64).reshape(1, -1) |
|
|
| past_kv = {} |
| for i in range(NUM_LAYERS): |
| past_kv[f"past_key_values.{i}.key"] = np.zeros( |
| (1, NUM_KV_HEADS, 0, HEAD_DIM), dtype=np.float32 |
| ) |
| past_kv[f"past_key_values.{i}.value"] = np.zeros( |
| (1, NUM_KV_HEADS, 0, HEAD_DIM), dtype=np.float32 |
| ) |
|
|
| eos_id = tok.eos_token_id |
| total_len = seq_len |
|
|
| for step in range(max_new_tokens): |
| feeds = { |
| "input_ids": input_ids, |
| "attention_mask": attention_mask, |
| "position_ids": position_ids, |
| **past_kv, |
| } |
| outputs = self.session.run(None, feeds) |
| logits = outputs[self.output_index["logits"]] |
| next_id = self._sample(logits[0, -1, :], temperature, top_p, top_k) |
|
|
| if next_id == eos_id: |
| break |
|
|
| piece = tok.decode([next_id], skip_special_tokens=True) |
| yield piece |
|
|
| input_ids = np.array([[next_id]], dtype=np.int64) |
| total_len += 1 |
| attention_mask = np.ones((1, total_len), dtype=np.int64) |
| position_ids = np.array([[total_len - 1]], dtype=np.int64) |
| past_kv = {} |
| for i in range(NUM_LAYERS): |
| past_kv[f"past_key_values.{i}.key"] = outputs[ |
| self.output_index[f"present.{i}.key"] |
| ] |
| past_kv[f"past_key_values.{i}.value"] = outputs[ |
| self.output_index[f"present.{i}.value"] |
| ] |
|
|
|
|
| |
| |
| |
| generator = Generator() |
| gen_lock = threading.Lock() |
|
|
|
|
| @app.on_event("startup") |
| def _startup(): |
| threading.Thread(target=generator.load, daemon=True).start() |
|
|
|
|
| |
| |
| |
| class Message(BaseModel): |
| role: str |
| content: str |
|
|
|
|
| class GenerateRequest(BaseModel): |
| messages: list[Message] |
| max_new_tokens: int = 600 |
| temperature: float = 0.7 |
| top_p: float = 0.8 |
| top_k: int = 20 |
|
|
|
|
| class HwpxRequest(BaseModel): |
| text: str |
|
|
|
|
| |
| |
| |
| STATIC_DIR = Path(__file__).parent / "static" |
| if STATIC_DIR.exists(): |
| app.mount("/static", StaticFiles(directory=str(STATIC_DIR)), name="static") |
|
|
|
|
| @app.get("/") |
| def root(): |
| index = STATIC_DIR / "index.html" |
| if index.exists(): |
| return FileResponse(str(index)) |
| return { |
| "status": generator.status, |
| "ready": generator.ready, |
| "model_repo": MODEL_REPO, |
| } |
|
|
|
|
| @app.get("/status") |
| def status(): |
| return { |
| "status": generator.status, |
| "ready": generator.ready, |
| "model_repo": MODEL_REPO, |
| } |
|
|
|
|
| @app.get("/health") |
| def health(): |
| return {"ready": generator.ready, "status": generator.status} |
|
|
|
|
| |
| |
| |
| def _ext_hwpx(path: Path) -> str: |
| HP = "http://www.hancom.co.kr/hwpml/2012/paragraph" |
| lines = [] |
| with zipfile.ZipFile(str(path), "r") as zf: |
| sections = sorted( |
| n for n in zf.namelist() |
| if n.startswith("Contents/section") and n.endswith(".xml") |
| ) |
| for name in sections: |
| with zf.open(name) as f: |
| try: |
| root = ET.parse(f).getroot() |
| except ET.ParseError: |
| continue |
| for para in root.iter(f"{{{HP}}}p"): |
| lines.append("".join( |
| t.text for t in para.iter(f"{{{HP}}}t") if t.text |
| )) |
| text = "\n".join(lines) |
| return re.sub(r"\n{3,}", "\n\n", text).strip() |
|
|
|
|
| def _ext_hwp(path: Path) -> str: |
| import struct, zlib |
| try: |
| import olefile |
| except ImportError: |
| raise RuntimeError("olefile 미설치 (pip install olefile)") |
|
|
| try: |
| ole = olefile.OleFileIO(str(path)) |
| except Exception as exc: |
| raise RuntimeError(f"HWP OLE 열기 실패 [{type(exc).__name__}]: {exc or 'no detail'}") |
|
|
| compressed = True |
| try: |
| if ole.exists("FileHeader"): |
| hdr = ole.openstream("FileHeader").read() |
| if len(hdr) >= 40: |
| compressed = bool(struct.unpack_from("<I", hdr, 36)[0] & 0x1) |
| except Exception: |
| pass |
|
|
| |
| try: |
| all_entries = ole.listdir() |
| print(f"[HWP] OLE entries: {all_entries[:20]}", flush=True) |
| sections = sorted( |
| (e[1] for e in all_entries |
| if isinstance(e, (list, tuple)) and len(e) >= 2 |
| and str(e[0]).lower() == "bodytext" |
| and str(e[1]).lower().startswith("section")), |
| key=lambda s: int(re.sub(r"\D", "", s) or "0"), |
| ) |
| except Exception as exc: |
| raise RuntimeError(f"섹션 목록 오류 [{type(exc).__name__}]: {exc or 'no detail'}") |
|
|
| if not sections: |
| entry_dump = str([str(e) for e in all_entries[:20]]) |
| raise RuntimeError(f"BodyText/Section* 없음. OLE 항목: {entry_dump}") |
|
|
| lines: list[str] = [] |
| for sec_name in sections: |
| try: |
| raw = ole.openstream(f"BodyText/{sec_name}").read() |
| except Exception as exc: |
| print(f"[HWP] openstream {sec_name} 오류: {exc}", flush=True) |
| continue |
|
|
| try: |
| data = zlib.decompress(raw, -15) if compressed else raw |
| except zlib.error: |
| try: |
| data = zlib.decompress(raw) |
| except zlib.error: |
| data = raw |
|
|
| pos = 0 |
| while pos + 4 <= len(data): |
| hval = struct.unpack_from("<I", data, pos)[0] |
| tag_id = hval & 0x3FF |
| size = (hval >> 20) & 0xFFF |
| pos += 4 |
| if size == 0xFFF: |
| if pos + 4 > len(data): break |
| size = struct.unpack_from("<I", data, pos)[0] |
| pos += 4 |
|
|
| if tag_id == 67 and size >= 2: |
| chunk = data[pos:pos + size] |
| t, chars = 0, [] |
| while t + 2 <= len(chunk): |
| cp = struct.unpack_from("<H", chunk, t)[0] |
| t += 2 |
| if cp == 0x0D: |
| lines.append("".join(chars)); chars = [] |
| elif 0x01 <= cp <= 0x0C: |
| t += 8 |
| elif cp >= 0x20: |
| chars.append(chr(cp)) |
| if chars: |
| lines.append("".join(chars)) |
| pos += size |
|
|
| return re.sub(r"\n{3,}", "\n\n", "\n".join(l for l in lines if l)).strip() |
|
|
|
|
| def _ext_docx(path: Path) -> str: |
| from docx import Document |
| doc = Document(str(path)) |
| return "\n".join(p.text for p in doc.paragraphs) |
|
|
|
|
| def _ext_pdf(path: Path) -> str: |
| from pypdf import PdfReader |
| reader = PdfReader(str(path)) |
| pages = [f"--- {i+1}페이지 ---\n{p.extract_text() or ''}" |
| for i, p in enumerate(reader.pages)] |
| return "\n".join(pages) |
|
|
|
|
| def _ext_txt(path: Path) -> str: |
| for enc in ["utf-8", "utf-8-sig", "cp949", "euc-kr"]: |
| try: |
| return path.read_text(encoding=enc) |
| except UnicodeDecodeError: |
| continue |
| raise RuntimeError("인코딩 감지 실패") |
|
|
|
|
| def _load_examples(doc_type: str, max_count: int = 3, max_chars: int = 2000) -> str: |
| folder = DATAS_DIR / doc_type |
| if not folder.exists(): |
| return "" |
| examples = [] |
| for p in sorted(folder.iterdir()): |
| if p.name.startswith(".") or p.stat().st_size == 0: |
| continue |
| try: |
| ext = p.suffix.lower() |
| if ext == ".hwpx": text = _ext_hwpx(p) |
| elif ext == ".hwp": text = _ext_hwp(p) |
| elif ext == ".docx": text = _ext_docx(p) |
| elif ext == ".pdf": text = _ext_pdf(p) |
| elif ext in (".txt", ".md"): text = _ext_txt(p) |
| else: continue |
| text = text.strip() |
| if text: |
| examples.append(f"--- 예시: {p.name} ---\n{text[:max_chars]}") |
| except Exception: |
| continue |
| if len(examples) >= max_count: |
| break |
| return "\n\n".join(examples) |
|
|
|
|
| @app.get("/examples/{doc_type}") |
| def get_examples(doc_type: str): |
| if doc_type not in DOC_TYPES: |
| raise HTTPException(400, f"알 수 없는 문서 유형: {doc_type}") |
| return {"doc_type": doc_type, "examples": _load_examples(doc_type)} |
|
|
|
|
| @app.post("/extract") |
| async def extract_file(file: UploadFile = File(...)): |
| """업로드된 문서(.hwp/.hwpx/.docx/.pdf/.txt)에서 텍스트를 추출해 반환.""" |
| ext = Path(file.filename or "").suffix.lower() |
| supported = {".hwp", ".hwpx", ".docx", ".pdf", ".txt"} |
| if ext not in supported: |
| raise HTTPException(400, f"지원하지 않는 형식: {ext}. 지원: {', '.join(sorted(supported))}") |
|
|
| content = await file.read() |
| with tempfile.NamedTemporaryFile(suffix=ext, delete=False) as tmp: |
| tmp.write(content) |
| tmp_path = Path(tmp.name) |
|
|
| try: |
| if ext == ".hwpx": |
| text = _ext_hwpx(tmp_path) |
| elif ext == ".hwp": |
| text = _ext_hwp(tmp_path) |
| elif ext == ".docx": |
| text = _ext_docx(tmp_path) |
| elif ext == ".pdf": |
| text = _ext_pdf(tmp_path) |
| else: |
| text = _ext_txt(tmp_path) |
| except Exception as e: |
| detail = str(e).strip() or type(e).__name__ |
| raise HTTPException(500, detail) |
| finally: |
| tmp_path.unlink(missing_ok=True) |
|
|
| return {"text": text, "filename": file.filename} |
|
|
|
|
| @app.post("/generate") |
| def generate(req: GenerateRequest): |
| """SSE 스트리밍 — 토큰별로 'data: {"token": "..."}\\n\\n' 전송.""" |
| if not generator.ready: |
| raise HTTPException(503, f"model not ready ({generator.status})") |
|
|
| if not gen_lock.acquire(blocking=False): |
| raise HTTPException(429, "another generation in progress (CPU 1요청 동시 처리)") |
|
|
| def event_stream(): |
| try: |
| t0 = time.time() |
| count = 0 |
| for piece in generator.stream( |
| [m.dict() for m in req.messages], |
| req.max_new_tokens, |
| req.temperature, |
| req.top_p, |
| req.top_k, |
| ): |
| count += 1 |
| yield f"data: {json.dumps({'token': piece}, ensure_ascii=False)}\n\n" |
| elapsed = time.time() - t0 |
| yield f"data: {json.dumps({'done': True, 'tokens': count, 'elapsed': round(elapsed, 1)}, ensure_ascii=False)}\n\n" |
| except Exception as e: |
| yield f"data: {json.dumps({'error': str(e)}, ensure_ascii=False)}\n\n" |
| finally: |
| gen_lock.release() |
|
|
| return StreamingResponse( |
| event_stream(), |
| media_type="text/event-stream", |
| headers={ |
| "Cache-Control": "no-cache", |
| "Connection": "keep-alive", |
| "X-Accel-Buffering": "no", |
| }, |
| ) |
|
|
|
|
| @app.post("/download/hwpx") |
| def download_hwpx(req: HwpxRequest): |
| """텍스트를 .hwpx 파일로 다운로드.""" |
| import re |
| from urllib.parse import quote |
| from fastapi.responses import Response |
|
|
| m = re.search(r"^제\s*목\s+(.+)$", req.text, re.MULTILINE) |
| title = m.group(1).strip() if m else "문서" |
| safe = re.sub(r'[\\/:*?"<>|]', "_", title)[:50] |
| encoded = quote(safe, safe="") |
|
|
| content = req.text.encode("utf-8-sig") |
|
|
| return Response( |
| content=content, |
| media_type="application/octet-stream", |
| headers={ |
| "Content-Disposition": f"attachment; filename*=UTF-8''{encoded}.hwpx", |
| }, |
| ) |
|
|
