Upload inference_example.py with huggingface_hub

inference_example.py

@@ -0,0 +1,130 @@
+"""
+HuggingFace에서 모델을 로드하여 추론하는 예제
+
+Usage:
+    from inference_example import extract_sentences
+    results = extract_sentences("삼성전자의 실적이 시장 예상을 상회했다. ...")
+"""
+
+import re
+from typing import List, Dict
+
+import torch
+from transformers import AutoTokenizer
+
+from model import (
+    DocumentEncoderConfig,
+    DocumentEncoderForExtractiveSummarization,
+    IDX_TO_ROLE,
+)
+
+
+def split_into_sentences(text: str) -> List[str]:
+    sentences = re.split(r"(?<=[.!?])\s+", text.strip())
+    return [s.strip() for s in sentences if s.strip()]
+
+
+def extract_sentences(
+    text: str,
+    model_name_or_path: str = "./",  # 로컬 또는 HuggingFace repo ID
+    top_k: int = 3,
+    threshold: float = 0.5,
+    device: str = None,
+) -> Dict:
+    """
+    텍스트에서 대표문장을 추출하고 역할을 분류합니다.
+
+    Args:
+        text: 입력 텍스트 (금융 리포트 등)
+        model_name_or_path: 모델 경로 또는 HuggingFace repo ID
+        top_k: 추출할 최대 문장 수
+        threshold: 대표문장 판단 임계값
+        device: cuda 또는 cpu
+
+    Returns:
+        dict with 'sentences', 'all_scores', 'all_roles', 'selected'
+    """
+    device = device or ("cuda" if torch.cuda.is_available() else "cpu")
+
+    config = DocumentEncoderConfig.from_pretrained(model_name_or_path)
+    model = DocumentEncoderForExtractiveSummarization.from_pretrained(
+        model_name_or_path, config=config
+    )
+    model = model.to(device)
+    model.eval()
+
+    tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)
+
+    sentences = split_into_sentences(text)
+    if not sentences:
+        return {"sentences": [], "all_scores": [], "all_roles": [], "selected": []}
+
+    max_sentences = config.max_sentences
+    max_length = config.max_length
+
+    padded = sentences[:max_sentences]
+    num_real = len(padded)
+    while len(padded) < max_sentences:
+        padded.append("")
+
+    all_input_ids, all_attention_mask = [], []
+    for s in padded:
+        if s:
+            enc = tokenizer(s, max_length=max_length, padding="max_length", truncation=True, return_tensors="pt")
+        else:
+            enc = {
+                "input_ids": torch.zeros(1, max_length, dtype=torch.long),
+                "attention_mask": torch.zeros(1, max_length, dtype=torch.long),
+            }
+        all_input_ids.append(enc["input_ids"])
+        all_attention_mask.append(enc["attention_mask"])
+
+    input_ids = torch.cat(all_input_ids, dim=0).unsqueeze(0).to(device)
+    attention_mask = torch.cat(all_attention_mask, dim=0).unsqueeze(0).to(device)
+    document_mask = torch.zeros(1, max_sentences, device=device)
+    document_mask[0, :num_real] = 1
+
+    with torch.no_grad():
+        scores, role_logits = model(input_ids, attention_mask, document_mask)
+
+    scores_list = scores[0, :num_real].tolist()
+    role_indices = role_logits[0, :num_real].argmax(dim=-1).tolist()
+    roles_list = [IDX_TO_ROLE[idx] for idx in role_indices]
+
+    selected = []
+    for i, (sent, score, role) in enumerate(zip(sentences, scores_list, roles_list)):
+        if score >= threshold:
+            selected.append({"index": i, "sentence": sent, "score": score, "role": role})
+
+    selected.sort(key=lambda x: x["score"], reverse=True)
+    selected = selected[:top_k]
+    selected.sort(key=lambda x: x["index"])
+
+    return {
+        "sentences": sentences,
+        "all_scores": scores_list,
+        "all_roles": roles_list,
+        "selected": selected,
+    }
+
+
+if __name__ == "__main__":
+    text = """
+    삼성전자의 2024년 4분기 실적이 시장 예상을 상회했다.
+    메모리 반도체 가격 상승으로 영업이익이 전분기 대비 30% 증가했다.
+    HBM3E 양산이 본격화되면서 AI 반도체 시장 점유율이 확대될 전망이다.
+    다만, 중국 시장의 불확실성이 여전히 리스크 요인으로 작용하고 있다.
+    회사는 올해 설비 투자를 20% 확대할 계획이다.
+    """
+
+    result = extract_sentences(text, model_name_or_path="./")
+
+    print("=" * 60)
+    print("전체 문장 분석:")
+    for i, (s, sc, r) in enumerate(zip(result["sentences"], result["all_scores"], result["all_roles"])):
+        marker = "*" if sc >= 0.5 else " "
+        print(f"  {marker} {i+1}. [{sc:.4f}] [{r:10s}] {s}")
+
+    print(f"\n선택된 대표문장:")
+    for item in result["selected"]:
+        print(f"  - [{item['score']:.4f}] [{item['role']:10s}] {item['sentence']}")