Update app.py

app.py

@@ -5,6 +5,11 @@ import seaborn as sns
 from transformers import pipeline
 import openai
 import os
+from sklearn.feature_extraction.text import TfidfVectorizer
+from sklearn.linear_model import LogisticRegression
+from sklearn.model_selection import train_test_split
+from sklearn.metrics import accuracy_score
+import joblib
 
 # Streamlit 페이지 설정을 가장 먼저 호출
 st.set_page_config(page_title="정치적 관점 분석", page_icon="📰", layout="wide")
@@ -54,10 +59,52 @@ def fetch_naver_news(query, display=5):
         st.error("뉴스 데이터를 불러오는 데 실패했습니다.")
         return []
 
-# 정치 성향 분석 모델 로드
-def load_sentiment_model():
-    classifier = pipeline("text-classification", model="bucketresearch/politicalBiasBERT")
-    return classifier
+# 머신러닝 모델 로드 및 학습
+def train_ml_model():
+    # 여기서는 샘플 데이터를 사용하여 학습
+    # 실제 데이터를 이용한 학습 과정이 필요합니다.
+    data = [
+        ("진보적인 정부 정책을 강화해야 한다", "LEFT"),
+        ("보수적인 경제 정책이 필요하다", "RIGHT"),
+        ("중립적인 입장에서 상황을 평가한다", "NEUTRAL")
+    ]
+    texts, labels = zip(*data)
+
+    # TF-IDF 벡터화
+    vectorizer = TfidfVectorizer(max_features=1000)
+    X = vectorizer.fit_transform(texts)
+    y = labels
+
+    # 훈련 및 테스트 데이터 나누기
+    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
+
+    # 로지스틱 회귀 모델 학습
+    model = LogisticRegression()
+    model.fit(X_train, y_train)
+
+    # 모델 성능 평가
+    y_pred = model.predict(X_test)
+    accuracy = accuracy_score(y_test, y_pred)
+    st.write(f"모델 정확도: {accuracy:.2f}")
+
+    # 모델 저장
+    joblib.dump(model, 'political_bias_model.pkl')
+    joblib.dump(vectorizer, 'tfidf_vectorizer.pkl')
+
+    return model, vectorizer
+
+# 로드된 머신러닝 모델로 성향 분석
+def analyze_article_sentiment_ml(text, model, vectorizer):
+    X = vectorizer.transform([text])
+    prediction = model.predict(X)[0]
+    
+    # 성향에 따른 레이블 반환
+    if prediction == "LEFT":
+        return "진보"
+    elif prediction == "RIGHT":
+        return "보수"
+    else:
+        return "중립"
 
 # GPT-4를 이용해 반대 관점 기사 생성
 def generate_article_gpt4(prompt):
@@ -76,28 +123,13 @@ def generate_article_gpt4(prompt):
     except Exception as e:
         return f"Error generating text: {e}"
 
-# 정치 성향 분석
-def analyze_article_sentiment(text, classifier):
-    result = classifier(text[:512])  # 너무 긴 텍스트는 잘라서 분석
-    label = result[0]["label"]
-    score = result[0]["score"]
-
-    # 모델에서 반환하는 라벨을 "진보", "보수", "중립"으로 매핑
-    if label == "LEFT":
-        return "진보", score
-    elif label == "RIGHT":
-        return "보수", score
-    else:
-        return "중립", score
-
 # 정치적 관점 비교 및 반대 관점 생성
-def analyze_news_political_viewpoint(query):
+def analyze_news_political_viewpoint(query, model, vectorizer):
     # 뉴스 데이터 가져오기
     news_items = fetch_naver_news(query)
     if not news_items:
         return [], {}
 
-    classifier = load_sentiment_model()
     results = []
     sentiment_counts = {"진보": 0, "보수": 0, "중립": 0}  # 매핑된 라벨에 맞게 초기화
 
@@ -107,8 +139,8 @@ def analyze_news_political_viewpoint(query):
         link = item["link"]  # 뉴스 링크 가져오기
         combined_text = f"{title}. {description}"
 
-        # 기사 성향 분석
-        sentiment, score = analyze_article_sentiment(combined_text, classifier)
+        # 머신러닝 모델을 이용한 성향 분석
+        sentiment = analyze_article_sentiment_ml(combined_text, model, vectorizer)
         sentiment_counts[sentiment] += 1  # 매핑된 키로 카운트 증가
 
         # 반대 관점 기사 생성
@@ -120,7 +152,6 @@ def analyze_news_political_viewpoint(query):
             "제목": title,
             "원본 기사": description,
             "성향": sentiment,
-            "성향 점수": score,
             "대조 관점 기사": opposite_article,
             "뉴스 링크": link  # 링크 추가
         })
@@ -146,29 +177,34 @@ def visualize_sentiment_distribution(sentiment_counts):
 st.title("📰 정치적 관점 비교 분석 도구")
 st.markdown("뉴스 기사의 정치 성향 분석과 반대 관점 기사를 생성하여 비교합니다.")
 
+# 머신러닝 모델 로드
+if not os.path.exists('political_bias_model.pkl'):
+    model, vectorizer = train_ml_model()
+else:
+    model = joblib.load('political_bias_model.pkl')
+    vectorizer = joblib.load('tfidf_vectorizer.pkl')
+
 # 사용자로부터 검색어 입력 받기
 query = st.text_input("검색 키워드를 입력하세요", value="정치")
 
 # 분석 시작 버튼
 if st.button("🔍 분석 시작"):
     with st.spinner("분석 중..."):
-        analysis_results, sentiment_counts = analyze_news_political_viewpoint(query)
+        analysis_results, sentiment_counts = analyze_news_political_viewpoint(query, model, vectorizer)
 
         if analysis_results:
             st.success("뉴스 분석이 완료되었습니다.")
             
-            # 성향 분포 시각화 (막대 그래프)
-            st.subheader("📊 성향 분포 시각화")
-            visualize_sentiment_distribution(sentiment_counts)
-            
-            # 상세 분석 결과 출력
-            st.subheader("📝 상세 분석 결과")
+            # 뉴스 기사 목록 표시
             for result in analysis_results:
-                st.write(f"#### {result['제목']}")
-                st.write(f"- **원본 기사**: {result['원본 기사']}")
-                st.write(f"- **성향**: {result['성향']} (점수: {result['성향 점수']:.2f})")
-                st.write(f"- **대조 관점 기사**: {result['대조 관점 기사']}")
-                st.write(f"- **뉴스 링크**: [링크]({result['뉴스 링크']})")  # 링크 출력
-                st.write("---")
+                st.subheader(result["제목"])
+                st.write(f"성향: {result['성향']}")
+                st.write(f"기사: {result['원본 기사']}")
+                st.write(f"[원본 기사 보기]({result['뉴스 링크']})")
+                st.write(f"대조 관점 기사: {result['대조 관점 기사']}")
+                st.markdown("---")
+            
+            # 성향 분포 시각화
+            visualize_sentiment_distribution(sentiment_counts)
         else:
-            st.error("분석된 뉴스 데이터가 없습니다.")
+            st.warning("검색된 뉴스가 없습니다.")