Update app.py

app.py

@@ -7,6 +7,7 @@ import matplotlib.pyplot as plt
 import csv
 import io
 import matplotlib.font_manager as fm
+from datetime import datetime, timedelta
 
 # 한국어 처리를 위한 KoSentence-BERT 모델 로드
 model = SentenceTransformer('jhgan/ko-sbert-sts')
@@ -14,15 +15,16 @@ model = SentenceTransformer('jhgan/ko-sbert-sts')
 # 나눔바른고딕 폰트 설정 (허깅페이스 환경에 맞게 수정)
 plt.rc('font', family='NanumBarunGothic')
 
-# 전역 변수로 분석 결과 저장
+# 전역 변수
 global_recommendations = None
 global_csv_string = None
+youtube_columns = None
 
 # CSV 문자열 생성 함수
 def create_csv_string(recommendations):
     output = io.StringIO()
     writer = csv.writer(output)
-    writer.writerow(["Employee ID", "Employee Name", "Recommended Programs"])
+    writer.writerow(["Employee ID", "Employee Name", "Recommended Programs", "Recommended YouTube Content"])
     for rec in recommendations:
         writer.writerow(rec)
     return output.getvalue()
@@ -71,12 +73,40 @@ def validate_and_get_columns(employee_df, program_df):
 
     return None, employee_cols, program_cols
 
+# 유튜브 데이터 열 선택 함수
+def select_youtube_columns(youtube_file):
+    global youtube_columns
+    youtube_df = pd.read_csv(youtube_file.name)
+    required_youtube_cols = ["title", "description", "url", "upload_date"]
+    youtube_columns = auto_match_columns(youtube_df, required_youtube_cols)
+    
+    column_options = {col: youtube_df.columns.tolist() for col in required_youtube_cols}
+    return gr.Dropdown.update(choices=youtube_df.columns.tolist(), value=youtube_columns.get("title")), \
+           gr.Dropdown.update(choices=youtube_df.columns.tolist(), value=youtube_columns.get("description")), \
+           gr.Dropdown.update(choices=youtube_df.columns.tolist(), value=youtube_columns.get("url")), \
+           gr.Dropdown.update(choices=youtube_df.columns.tolist(), value=youtube_columns.get("upload_date"))
+
+# 유튜브 콘텐츠 데이터 로드 및 처리 함수
+def load_youtube_content(file_path, title_col, description_col, url_col, upload_date_col):
+    youtube_df = pd.read_csv(file_path)
+    youtube_df = youtube_df[[title_col, description_col, url_col, upload_date_col]]
+    youtube_df.columns = ['title', 'description', 'url', 'upload_date']
+    youtube_df['upload_date'] = pd.to_datetime(youtube_df['upload_date'])
+    return youtube_df
+
+# 유튜브 콘텐츠와 교육 프로그램 매칭 함수
+def match_youtube_content(program_skills, youtube_df, model):
+    youtube_embeddings = model.encode(youtube_df['description'].tolist())
+    program_embeddings = model.encode(program_skills)
+    similarities = cosine_similarity(program_embeddings, youtube_embeddings)
+    return similarities
+
 # 직원 데이터를 분석하여 교육 프로그램을 추천하고, 테이블과 그래프를 생성하는 함수
-def hybrid_rag(employee_file, program_file):
+def hybrid_rag(employee_file, program_file, youtube_file, title_col, description_col, url_col, upload_date_col):
     global global_recommendations
     global global_csv_string
     
-    # 1. VectorRAG: KoSentence-BERT를 이용한 유사도 계산
+    # 직원 및 프로그램 데이터 로드
     employee_df = pd.read_csv(employee_file.name)
     program_df = pd.read_csv(program_file.name)
 
@@ -91,22 +121,38 @@ def hybrid_rag(employee_file, program_file):
 
     similarities = cosine_similarity(employee_embeddings, program_embeddings)
 
+    # 유튜브 콘텐츠 로드 및 처리
+    youtube_df = load_youtube_content(youtube_file.name, title_col, description_col, url_col, upload_date_col)
+    
+    # 유튜브 콘텐츠와 교육 프로그램 매칭
+    youtube_similarities = match_youtube_content(program_df[program_cols['skills_acquired']].tolist(), youtube_df, model)
+    
     recommendations = []
-    recommendation_rows = []  # 테이블 및 CSV로 저장할 데이터
+    recommendation_rows = []
     for i, employee in employee_df.iterrows():
         recommended_programs = []
+        recommended_youtube = []
         for j, program in program_df.iterrows():
             if similarities[i][j] > 0.5:
                 recommended_programs.append(f"{program[program_cols['program_name']]} ({program[program_cols['duration']]})")
+                
+                # 해당 프로그램과 가장 유사한 유튜브 콘텐츠 찾기
+                top_youtube_indices = youtube_similarities[j].argsort()[-3:][::-1]  # 상위 3개
+                for idx in top_youtube_indices:
+                    recommended_youtube.append(f"{youtube_df.iloc[idx]['title']} (URL: {youtube_df.iloc[idx]['url']})")
         
         if recommended_programs:
             recommendation = f"직원 {employee[employee_cols['employee_name']]}의 추천 프로그램: {', '.join(recommended_programs)}"
-            recommendation_rows.append([employee[employee_cols['employee_id']], employee[employee_cols['employee_name']], ", ".join(recommended_programs)])
+            youtube_recommendation = f"추천 유튜브 콘텐츠: {', '.join(recommended_youtube)}"
+            recommendation_rows.append([employee[employee_cols['employee_id']], employee[employee_cols['employee_name']], 
+                                        ", ".join(recommended_programs), ", ".join(recommended_youtube)])
         else:
             recommendation = f"직원 {employee[employee_cols['employee_name']]}에게 적합한 프로그램이 없습니다."
-            recommendation_rows.append([employee[employee_cols['employee_id']], employee[employee_cols['employee_name']], "적합한 프로그램 없음"])
+            youtube_recommendation = "추천할 유튜브 콘텐츠가 없습니다."
+            recommendation_rows.append([employee[employee_cols['employee_id']], employee[employee_cols['employee_name']], 
+                                        "적합한 프로그램 없음", "추천 콘텐츠 없음"])
 
-        recommendations.append(recommendation)
+        recommendations.append(recommendation + "\n" + youtube_recommendation)
 
     global_recommendations = recommendation_rows
 
@@ -129,7 +175,7 @@ def hybrid_rag(employee_file, program_file):
     global_csv_string = create_csv_string(recommendation_rows)
     
     # 결과 테이블 데이터프레임 생성
-    result_df = pd.DataFrame(recommendation_rows, columns=["Employee ID", "Employee Name", "Recommended Programs"])
+    result_df = pd.DataFrame(recommendation_rows, columns=["Employee ID", "Employee Name", "Recommended Programs", "Recommended YouTube Content"])
     
     return result_df, chart_buffer, gr.File.update(visible=True)
 
@@ -141,7 +187,7 @@ def chat_response(message, history):
     
     for employee in global_recommendations:
         if employee[1].lower() in message.lower():
-            return f"{employee[1]}님에게 추천된 프로그램은 다음과 같습니다: {employee[2]}"
+            return f"{employee[1]}님에게 추천된 프로그램은 다음과 같습니다: {employee[2]}\n\n추천 유튜브 콘텐츠: {employee[3]}"
     
     return "죄송합니다. 해당 직원의 정보를 찾을 수 없습니다. 다른 직원 이름을 입력해주세요."
 
@@ -154,13 +200,23 @@ def download_csv():
 
 # Gradio 블록
 with gr.Blocks(css=".gradio-button {background-color: #007bff; color: white;} .gradio-textbox {border-color: #6c757d;}") as demo:
-    gr.Markdown("<h1 style='text-align: center; color: #2c3e50;'>💼 HybridRAG 시스템</h1>")
+    gr.Markdown("<h1 style='text-align: center; color: #2c3e50;'>💼 HybridRAG 시스템 (유튜브 콘텐츠 포함)</h1>")
     
     with gr.Row():
         with gr.Column(scale=1, min_width=300):
-            gr.Markdown("<h3 style='color: #34495e;'>1. 직원 및 프로그램 데이터를 업로드하세요</h3>")
+            gr.Markdown("<h3 style='color: #34495e;'>1. 데이터를 업로드하세요</h3>")
             employee_file = gr.File(label="직원 데이터 업로드", interactive=True)
             program_file = gr.File(label="교육 프로그램 데이터 업로드", interactive=True)
+            youtube_file = gr.File(label="유튜브 콘텐츠 데이터 업로드", interactive=True)
+            
+            gr.Markdown("<h4 style='color: #34495e;'>유튜브 데이터 열 선택</h4>")
+            title_col = gr.Dropdown(label="제목 열")
+            description_col = gr.Dropdown(label="설명 열")
+            url_col = gr.Dropdown(label="URL 열")
+            upload_date_col = gr.Dropdown(label="업로드 날짜 열")
+            
+            youtube_file.change(select_youtube_columns, inputs=[youtube_file], outputs=[title_col, description_col, url_col, upload_date_col])
+            
             analyze_button = gr.Button("분석 시작", elem_classes="gradio-button")
             output_table = gr.DataFrame(label="분석 결과 (테이블)")
             csv_download = gr.File(label="추천 결과 다운로드", visible=False)
@@ -169,13 +225,15 @@ with gr.Blocks(css=".gradio-button {background-color: #007bff; color: white;} .g
             gr.Markdown("<h3 style='color: #34495e;'>2. 분석 결과 및 시각화</h3>")
             chart_output = gr.Image(label="시각화 차트")
 
-    gr.Markdown("<h3 style='color: #34495e;'>3. 직원별 추천 프로그램 확인</h3>")
+    gr.Markdown("<h3 style='color: #34495e;'>3. 직원별 추천 프로그램 및 유튜브 콘텐츠 확인</h3>")
     chatbot = gr.Chatbot()
     msg = gr.Textbox(label="직원 이름을 입력하세요")
     clear = gr.Button("대화 내역 지우기")
 
     # 분석 버튼 클릭 시 테이블, 차트, 파일 다운로드를 업데이트
-    analyze_button.click(hybrid_rag, inputs=[employee_file, program_file], outputs=[output_table, chart_output, csv_download])
+    analyze_button.click(hybrid_rag, 
+                         inputs=[employee_file, program_file, youtube_file, title_col, description_col, url_col, upload_date_col], 
+                         outputs=[output_table, chart_output, csv_download])
 
     # CSV 다운로드 버튼
     csv_download.click(download_csv, inputs=[], outputs=[csv_download])