|
| 1 | +import re |
| 2 | +from sklearn.metrics.pairwise import cosine_similarity |
| 3 | +from typing import List, Dict |
| 4 | +import numpy as np |
| 5 | +from app.models.skt_kobert import get_sentence_embedding |
| 6 | + |
| 7 | + |
| 8 | +CATEGORY_LABELS = { |
| 9 | + "여행": ["도쿄 여행", "유럽 배낭 여행", "국내 캠핑", "비행기 예약", "숙소 추천"], |
| 10 | + "교통": ["지하철 환승", "고속열차 이용", "비행기 탑승", "대중교통 이용", "렌터카 예약"], |
| 11 | + "쇼핑": ["면세점 할인", "백화점 세일", "전자제품 구매", "패션 브랜드 쇼핑", "기념품 구매"], |
| 12 | + "음식": ["일본 라멘 맛집", "프랑스 빵 추천", "한식당 방문", "카페 탐방", "길거리 음식"], |
| 13 | + "기타": ["문화 체험", "미술관 방문", "박물관 견학", "자연 탐방", "테마파크 방문"] |
| 14 | +} |
| 15 | + |
| 16 | +url_pattern = re.compile(r'https?://[a-zA-Z0-9./?=&_%:-]+') |
| 17 | + |
| 18 | +# ✅ 문장에서 URL을 추출하고 문장 내에서 분리하는 함수 |
| 19 | +def extract_urls_from_sentences(sentence: str): |
| 20 | + urls = url_pattern.findall(sentence) |
| 21 | + text_without_urls = url_pattern.sub(' ', sentence).strip() |
| 22 | + text_without_urls = re.sub(r'\s+', ' ', text_without_urls) |
| 23 | + return text_without_urls, urls |
| 24 | + |
| 25 | +category_embeddings = { |
| 26 | + category: np.mean([get_sentence_embedding(example) for example in examples], axis=0) |
| 27 | + for category, examples in CATEGORY_LABELS.items() |
| 28 | +} |
| 29 | + |
| 30 | +def generate_new_category(text: str): |
| 31 | + words = text.split() |
| 32 | + for word in words: |
| 33 | + if len(word) > 1: |
| 34 | + return word |
| 35 | + return "기타" |
| 36 | + |
| 37 | +def classify_paragraph(paragraph: str, threshold: float = 0.7): |
| 38 | + global category_embeddings |
| 39 | + sentences = paragraph.split("\n") |
| 40 | + processed_sentences = [] |
| 41 | + paragraph_embedding = get_sentence_embedding(paragraph) |
| 42 | + |
| 43 | + |
| 44 | + best_category = None |
| 45 | + best_similarity = 0 |
| 46 | + |
| 47 | + for category, category_vector in category_embeddings.items(): |
| 48 | + similarity = cosine_similarity([paragraph_embedding], [category_vector])[0][0] |
| 49 | + if similarity > best_similarity: |
| 50 | + best_similarity = similarity |
| 51 | + best_category = category |
| 52 | + |
| 53 | + # ✅ 기존 카테고리에 없으면 "category": "no", 추천 카테고리 제공 |
| 54 | + if best_similarity < threshold: |
| 55 | + recommend_category = generate_new_category(paragraph) |
| 56 | + if recommend_category not in category_embeddings: |
| 57 | + category_embeddings[recommend_category] = paragraph_embedding |
| 58 | + return_category = "no" |
| 59 | + else: |
| 60 | + recommend_category = best_category |
| 61 | + return_category = best_category |
| 62 | + |
| 63 | + |
| 64 | + for sentence in sentences: |
| 65 | + text_part, urls = extract_urls_from_sentences(sentence) |
| 66 | + sentence_embedding = get_sentence_embedding(text_part) if text_part else None |
| 67 | + |
| 68 | + |
| 69 | + if urls: |
| 70 | + sub_category = "관련 링크" |
| 71 | + else: |
| 72 | + sub_category = recommend_category |
| 73 | + |
| 74 | + processed_sentences.append({ |
| 75 | + "text": text_part if text_part else "URL 포함 문장", |
| 76 | + "sub_category": sub_category, |
| 77 | + "urls": urls if urls else None |
| 78 | + }) |
| 79 | + |
| 80 | + return { |
| 81 | + "category": return_category, |
| 82 | + "recommend_category": recommend_category, |
| 83 | + "sentences": processed_sentences |
| 84 | + } |
0 commit comments