work / app.py
eunbi6768's picture
Upload 4 files
3901386 verified
import streamlit as st
import numpy as np
import faiss
import pickle
from sentence_transformers import SentenceTransformer
import pandas as pd
# 제목 표시
st.title('직장 괴롭힘 참고 사례 검색과 판단')
# SentenceTransformer 모델 로드
model = SentenceTransformer('bongsoo/kpf-sbert-v1') # 768차원
# 참고 사례 데이터프레임 로드
df = pd.read_excel('직장괴롭힘_071424.xlsx')
# 문장 임베딩
#embeddings = model.encode(df['사례'].tolist())
# 저장된 임베딩을 pickle 파일에서 로드
with open('embeddingsNorm_직장괴롭힘071424_kpfsbert.pkl', 'rb') as f:
embeddings = pickle.load(f)
# 임베딩을 코사인 유사도 계산을 위해 정규화
embeddings_norm = embeddings / np.linalg.norm(embeddings, axis=1, keepdims=True)
# FAISS 인덱스 생성 (내적 검색용)
d = embeddings_norm.shape[1] # 임베딩의 차원 수를 저장: embeddings_norm은 (문장 수, 차원 수) 형태의 2차원 배열이므로, shape[1]은 임베딩의 차원 수를 나타냄
index = faiss.IndexFlatIP(d) # 내적(IP, Inner Product) 검색을 위한 FAISS 인덱스 생성. IndexFlatIP는 내적 기반의 유사도 검색을 위한 인덱스를 의미
index.add(embeddings_norm) # 정규화된 임베딩을 인덱스에 추가. 검색을 위해 미리 계산된 모든 문장 임베딩을 인덱스에 추가하여 검색이 가능하도록 함
# 찾고자 하는 문장을 입력받음
query_sentence = st.text_area('검색할 문장을 입력하세요', '''예시: 회사 과장님이 회의에서 "조용히 있으라"며 야단치고, 동료들도 회식 등 모임에서 따돌림합니다''')
if st.button('분석'):
# 입력된 문장의 임베딩
query_embedding = model.encode([query_sentence])
# 입력된 문장의 임베딩을 정규화
query_embedding_norm = query_embedding / np.linalg.norm(query_embedding, axis=1, keepdims=True)
# 유사도 검색 (상위 5개 문장)
k = 5
distances, indices = index.search(query_embedding_norm, k)
# 내적 값을 코사인 유사도로 변환
cosine_similarities = distances[0]
# 인덱스와 유사도를 저장하는 딕셔너리 생성
#indices_distances = {'indices': indices[0], 'cosine_similarities': cosine_similarities}
# 결과를 저장할 데이터프레임 생성
df_result = df.iloc[indices[0]].copy()
df_result['유사도'] = cosine_similarities
# 유사도 기준 내림차순으로 정렬
df_result = df_result.sort_values(by='유사도', ascending=False)
# 결과 표시
st.write('분석 결과:')
st.write(df_result)