app.py

# Import custom modules for prediction
from predictor import predict
import google.generativeai as genai
import streamlit as st
import pandas as pd
import wikipedia
import json
import regex as re
from typing import List, Dict, Any
import google.generativeai as genai
import warnings
import logging
import time

# Disable warnings and logging
warnings.filterwarnings('ignore')
logging.disable(logging.WARNING)
def split_sentences_vietnamese(text):
    # Thay thế dấu ba chấm tạm thời
    text = text.replace("...", " @@@ ")

    # Chia câu bằng các dấu chấm, chấm hỏi, chấm than
    sentences = re.split(r'(?<!\w\.\w.)(?<!\w\.\w)(?<!\w\.)[.!?]\s+', text)
    
    # Khôi phục dấu ba chấm và làm sạch dữ liệu
    sentences = [s.replace(" @@@ ", "...").strip() for s in sentences if s.strip()]

    return sentences
# Set Wikipedia language to Vietnamese
wikipedia.set_lang("vi")

# Character to number mapping for Vietnamese alphabet
char_to_num = {
    'A': 1, 'Ă': 2, 'Â': 3, 'B': 4, 'C': 5, 'D': 6, 'Đ': 7, 'E': 8, 'Ê': 9, 'G': 10,
    'H': 11, 'I': 12, 'K': 13, 'L': 14, 'M': 15, 'N': 16, 'O': 17, 'Ô': 18, 'Ơ': 19,
    'P': 20, 'Q': 21, 'R': 22, 'S': 23, 'T': 24, 'U': 25, 'Ư': 26, 'V': 27, 'X': 28, 'Y': 29
}

class ClaimGenerator:
    """Class to handle advanced claim generation with different APIs."""
    
    def __init__(self, claim_type: str, api_keys_file: str = "API.csv"):
        self.claim_type = claim_type
        self.api_keys = self._load_api_keys(api_keys_file)
        self.current_key_index = 0
        self._configure_api()
        API_KEYS = {
            "SUPPORTED": "AIzaSyCWoLgqVB8WLM7JPQMz4u3_jIOsHVs_SkA",
            "REFUTED": "AIzaSyBe4ycn90utUlwhbD-fiyHJCitsPqVUO20",
            "NEI": "AIzaSyBqcQXJlF18yItimSQRtVcKCSpxguLa360"
        }
    
        self.SCHEMAS = {
            "SUPPORTED": {
                "type": "object",
                "properties": {
                    "claims": {
                        "type": "array",
                        "items": {
                            "type": "string"
                        },
                        "description": """
                        Câu nhận định đúng với các thông tin và số liệu được trích xuất trực tiếp từ context,yêu cầu phức tạp và dễ bị nhầm lẫn thành nhãn REFUTED, kết hợp nhiều kỹ thuật biến đổi:
                        Yêu cầu:
                        - Thông tin phải chính xác 100% so với context
                        - Không được thêm thông tin ngoài context
                        - Có thể kết hợp nhiều thông tin từ các phần khác nhau của context
                        - Phải có thể verify được từ context
                        Gợi ý cách đặt câu:
                        - Thay thế các bằng các từ đồng nghĩa
                        - Thay thế các lượng từ có cùng ý nghĩa
                        - Chuyển từ chủ động thành bị động & ngược lại
                        - Dùng từ trái nghĩa + phủ định
                        - Lược bỏ những bổ ngữ "không cần thiết" (không làm thay đổi nội dung của câu)
                        - Chuyển đổi danh từ, tân ngữ thành 1 mệnh đề quan hệ
                        - Tạo câu tiền giả định
                        Ví dụ:
                        VD1:
                        Claim: Quân Hán của Hàn Tín đã hủy diệt quân Triệu ở trận Tình Hình.
                        Context: ^  Năm 204 TCN, Trận Tỉnh Hình (trận Bối Thủy - Hàn Tín phá Triệu), quân nước Hán do danh tướng Hàn Tín chỉ huy đập tan nát quân Triệu của tướng Trần Dư. 
                        VD2:
                        Claim: Trong quá trình làm việc nếu có tương tác với những vật dụng dính virus thì đều có nguy cơ nhiễm bệnh.
                        Context: Nhân viên Y tế: Bất kỳ nhân viên Y tế nào (Bác sĩ, Nha sỹ, Y tá, Y công.) cũng có thể bị nhiễm Viêm gan siêu vi C do tiếp xúc với bệnh phẩm chứa siêu vi trong quá trình làm việc. 
                        VD3:
                        Claim: Trần Bảo dẫn quân sang cứu viện phủ thành Đại La nhưng thất bại.
                        Context: Năm 931, Dương Đình Nghệ là tướng cũ của Khúc Hạo đem quân đánh phủ thành Đại La, lại đánh tan quân Nam Hán do Trần Bảo dẫn sang cứu viện, giết Trần Bảo và tự xưng là Tiết độ sứ. 
                        VD4:
                        Claim: Gan bị xơ sẽ không dễ dàng khôi phục lại.
                        Context: Khi đã xơ, gan khó hồi phục lại, cho dù tình trạng viêm có thuyên giảm. Vì vậy, các thầy thuốc khuyên nên điều trị sớm nhằm ngăn ngừa hoặc làm chậm tiến triển sang giai đoạn xơ gan. 
                        VD5:
                        Claim: Nền văn minh này đóng góp một trong những văn bản lâu đời nhất thế giới.
                        Context: Các văn bản lịch sử văn hóa đã góp phần vào sự tồn tại của nền văn minh này. Một trong những văn bản cổ nhất trên thế giới, có tên Sử thi Gilgamesh, có nguồn gốc từ nền văn minh này. 
                        VD6:
                        Claim: Sau khi thua trận trước quân Hán, Hai Bà Trưng đã tuẫn tiết.
                        Context: Năm 43, nhà Hán sai Phục ba tướng quân Mã Viện đem quân sang tái chiếm và đàn áp, Hai Bà Trưng chống không nổi phải rút về giữ Cấm Khê rồi tự vẫn ở sông Hát. Dân địa phương lập đền thờ ở Hát Giang. 
                        """
                    }
                },
                "required": ["claims"]
            },
            "REFUTED": {
                "type": "object",
                "properties": {
                    "claims": {
                        "type": "array",
                        "items": {
                            "type": "string"
                        },
                        "description": """
                        Câu nhận định sai, mâu thuẫn với thông tin trong context, tinh vi, khó phát hiện và dễ bị nhầm lẫn thành nhãn SUPPORTED:
                        Yêu cầu:
                        - Phải sử dụng cùng chủ đề và phạm vi với context
                        - Thông tin hoặc số liệu phải mâu thuẫn rõ ràng với context
                        - Có thể thay đổi số liệu, sự kiện, hoặc mối quan hệ
                        - Phải có thể verify được là sai dựa trên context
                        Gợi ý cách đặt câu:
                        - Xáo trộn các giải thích của các đối tượng với nhau
                        - Thay thế bằng các từ trái nghĩa
                        - Thay thế các từ ngữ trong cùng lĩnh vực
                        - Biến đổi các thông số, đơn vị của các đối tượng
                        Ví dụ:
                        VD1:
                        Claim: Các phụ âm của phương ngữ California của tiếng Anh Mỹ, và của chuẩn RP được mô tả qua bảng phía trên.
                        Context: Đa số phương ngữ tiếng Anh sở hữu cùng 24 phụ âm giống nhau. Bảng bên dưới thể hiện các phụ âm của phương ngữ California của tiếng Anh Mỹ, và của chuẩn RP. 
                        VD2:
                        Claim: Mặc dù đã xảy ra sự kiện đắm tàu Amphitrite vào năm 1698 nhưng quần đảo vẫn không một ai biết đến và ngày càng bị lu mờ đi.
                        Context: Năm 1698: Quần đảo trở nên nổi tiếng trong các biên niên sử hàng hải với vụ đắm tàu Amphitrite dưới thời vua Louis XIV trong khi đi từ Pháp sang Trung Quốc. 
                        VD3:
                        Claim: Lí do để quần đảo trở nên nổi tiếng và được nhiều người biết đến trong các biên niên sử của hàng hải là vụ đắm tàu Amphitrite vào năm 1700.
                        Context: Năm 1698: Quần đảo trở nên nổi tiếng trong các biên niên sử hàng hải với vụ đắm tàu Amphitrite dưới thời vua Louis XIV trong khi đi từ Pháp sang Trung Quốc. 
                        VD4:
                        Claim: Trôi dạt lục địa đã đẩy toàn bộ bề mặt đất liền lại với nhau.
                        Context: Trôi dạt lục địa đã tái định hình thể cho các đại dương của Trái Đất, kết hợp và chia cắt các đại dương cổ để tạo ra các đại dương như hiện nay. Các đại dương cổ có: 
                        VD5:
                        Claim: Các tài liệu ghi về thời kỳ bắc thuộc cuối cùng của Đại Việt đã biến mất hoàn toàn từ xưa cho tới nay.
                        Context: Annam Lưu trữ 2006-12-03 tại Wayback Machine, thu thập các đoạn tài liệu Minh thực lục (明實錄) của nhà Minh về Đại Việt, trong đó có phần về thời kỳ bắc thuộc lần cuối. 
                        VD6
                        Claim: Hiện tượng mô não mất quyền kiểm soát trong cơ thể là vì bệnh thiếu máu cục bộ khiến máu tràn một lượng lớn vào não.
                        Context: Trong đột quỵ do thiếu máu cục bộ, lượng máu cung cấp cho một phần não bị giảm, dẫn đến rối loạn chức năng của mô não ở khu vực đó. Có bốn lý do tại sao điều này có thể xảy ra: 
                        """
                    }
                },
                "required": ["claims"]
            },
            "NEI": {
                "type": "object",
                "properties": {
                    "claims": {
                        "type": "array",
                        "items": {
                            "type": "string"
                        },
                        "description": """
                        Câu có thông tin không thể xác minh tính chính xác chỉ dựa vào context, có khả năng gây nhầm lẫn với SUPPORTED hoặc REFUTED:
                        
                        Yêu cầu cơ bản:
                        - Phải liên quan chặt chẽ đến chủ đề chính của context
                        - Có ít nhất 1 thông tin không được có trong context
                        - Không được mâu thuẫn trực tiếp với context
                        - Không thể xác định được đúng sai chỉ dựa vào context
                        - Phải tạo ra sự mơ hồ và không chắc chắn cao
                         Gợi ý cách đặt câu:
                        - Tạo câu một phần thông tin liên quan đến context, nhưng đã thêm vào những chi tiết, suy luận hoặc kết luận không được hỗ trợ trực tiếp bởi context.
                        Ví dụ:
                        VD1:
                        Claim: Thời kỳ bắc thuộc lần cuối không được biết đến cho tới khi Annam lưu trữ xuất hiện.
                        Context: Annam Lưu trữ 2006-12-03 tại Wayback Machine, thu thập các đoạn tài liệu Minh thực lục (明實錄) của nhà Minh về Đại Việt, trong đó có phần về thời kỳ bắc thuộc lần cuối. 
                        VD2:
                        Claim: Phát triển du lịch là mục tiêu phấn đấu của tỉnh Nghệ An.
                        Context: Nghệ An còn lưu giữ được nhiều di tích văn hoá lịch sử, nhiều danh lam thắng cảnh, lễ hội văn hoá truyền thống - đó là những yếu tố thuận lợi giúp cho du lịch Nghệ Antriển. 
                        VD3:
                        Claim: Quân đội Ba Tư đông hơn nhưng vẫn không địch lại người La Mã.
                        Context: ^  Năm 627, Trận Nineveh (trong cuộc Chiến tranh La Mã-Ba Tư), quân Đông La Mã do đích thân Hoàng đế Heraclius thống suất đánh thắng quân Ba Tư do Đại tướng Rhahzadh cầm đầu. 
                        VD4:
                        Claim: Bệnh phẩm chứa siêu vi viêm gan C được đề xuất nên bị loại bỏ hoàn toàn.
                        Context: Nhân viên Y tế: Bất kỳ nhân viên Y tế nào (Bác sĩ, Nha sỹ, Y tá, Y công.) cũng có thể bị nhiễm Viêm gan siêu vi C do tiếp xúc với bệnh phẩm chứa siêu vi trong quá trình làm việc. 
                        VD5:
                        Claim: Thập Nhị Chi là tính toán thời gian theo 12 con vật.
                        Context: Tương truyền ngày xưa có một người tên Đại Nhiêu đã lập ra Thập Can và Thập Nhị Chi để giúp người ta tính toán thời gian. Việc tính giờ cũng có liên quan đến tập tính của 12 loài vật: 
                        VD5:
                        Claim: Nghệ An là một tỉnh nằm ở phía bắc sông Lam.
                        Context: Năm 1831, vua Minh Mệnh chia trấn Nghệ An thành 2 tỉnh: Nghệ An (phía Bắc sông Lam); Hà Tĩnh (phía nam sông Lam). Sau đó hai tỉnh Nghệ An và Hà Tĩnh sáp nhập lại, lấy tên là tỉnh An Tĩnh. 
                        VD6:
                        Claim: Được mô phỏng giống như lễ hội đường phố tại các nước Châu Âu, Carneval Hội An đã được tổ chức tại Hội An vào Giao thừa 2009
                        Context: Carneval Hội An là lễ hội đường phố được tổ chức lần đầu tiên tại thành phố Hội An vào Giao thừa năm 2009 (dương lịch). Lễ hội mô phỏng theo các lễ hội Carneval đường phố vốn rất nổi tiếng tại các nước châu Âu và Mỹ Latin .
                        Lưu ý quan trọng:
                        - Ưu tiên sử dụng các từ ngữ gợi ý về quan hệ nhân quả, so sánh, hoặc khái quát hóa
                        - Tránh tạo ra các claim quá rõ ràng là không có trong context
                        - Nên tận dụng các thông tin ngữ cảnh để tạo ra các suy luận hợp lý nhưng không thể kiểm chứng
                        - Kết hợp nhiều kỹ thuật để tăng độ khó trong việc xác định nhãn
                        - Đảm bảo câu claim vẫn giữ được tính logic và khả năng xảy ra trong thực tế
                        """
                    }
                },
                "required": ["claims"]
            }
        }
    
    def _load_api_keys(self, file_path: str ) -> List[str]:
        """Load API keys from CSV file for the specific claim type."""
        try:
            df = pd.read_csv(file_path)
            # Get all keys for the specific claim type, excluding NaN values
            keys = df[self.claim_type].dropna().tolist()
            if not keys:
                raise ValueError(f"No API keys found for {self.claim_type}")
            return keys
        except Exception as e:
            raise Exception(f"Error loading API keys: {str(e)}")

    def _configure_api(self):
        """Configure the API with the current key."""
        genai.configure(api_key=self.api_keys[self.current_key_index])

    def _rotate_api_key(self):
        """Rotate to the next available API key."""
        self.current_key_index = (self.current_key_index + 1) % len(self.api_keys)
        self._configure_api()
        print(f"Rotating to next API key for {self.claim_type}")


    def generate_claims(self, context: str, max_retries: int = 10) -> List[str]:
        """Generate claims with automatic API key rotation on errors."""
        retries = 0
        while retries < max_retries:
            try:
                model = genai.GenerativeModel(
                    'models/gemini-1.5-flash-latest',
                    generation_config={
                        "response_mime_type": "application/json",
                        "response_schema": self.SCHEMAS[self.claim_type],
                        "temperature": 0.9,
                        "top_p": 0.95,
                        "top_k": 40,
                        "max_output_tokens": 1024
                    }
                )

                prompt = f"""
                    Fact checking là 1 trong lĩnh vực xử lý ngôn ngữ tự nhiên (NLP), là bài toán kiểm tra và xác minh tính chính xác của 1 nhận định được đưa ra trên các thông tin, tin tức từ mạng xã hội, báo chí. Mô hình NLP sẽ xác định câu nhận định đó là chính xác hay sai sự thật hoặc không được đề cập đến trong thông tin đã được cung cấp trước đó (context).
               	    Trong bối cảnh của Adversarial, đây là bài toán tạo ra bộ dữ liệu đối kháng với mục đích đánh lừa gây và gây khó khăn cho các mô hình NLP để dự đoán tính xác thực của các câu nhận định đó. Mục đích của bài toán này là thử thách khả năng suy luận của các mô hình NLP, từ đó để tạo nên bộ dữ liệu mới chất lượng hơn, cải thiện “khả năng học” của mô hình học máy.
                    Context: là nguồn thông tin, nội dung được cung cấp hiện có trong dataset mà annotator sẽ dựa vào để tiến hành đặt câu và đánh dấu.
                    Câu Claim: là câu nhận định, tuyên bố viết về bất kỳ chủ đề gì, được xác định tính đúng sai hoặc không liên quan dựa trên thông tin từ context đưa ra mà không sử dụng bất kỳ kiến thức hoặc thông tin ngoài lề
                    + SUPPORTED là nhãn dành cho câu nhận định đúng, chính xác với các thông tin, số liệu trong câu được trích xuất từ đoạn văn gốc
                    + REFUTED là nhãn dành cho câu nhận định sai, với các thông tin, số liệu trong câu bị sai so với đoạn văn gốc
                    + NEI là nhãn dành cho câu có các thông tin, số liệu chưa được xác định là đúng hay sai dựa trên đoạn văn gốc
                    Context: "{context}"
                    Nhiệm vụ: Tạo một câu claim {self.claim_type} có liên quan đến context có độ khó cao, có thể gây khó khăn cho các mô hình NLP trong việc dự đoán nhãn chính xác.
                    Yêu cầu chi tiết:
                    1. Tuân thủ chặt chẽ các yêu cầu và kỹ thuật được mô tả trong schema
                    2. Tạo câu có cấu trúc phức tạp, kết hợp nhiều kỹ thuật biến đổi
                    3. Đảm bảo câu vẫn hợp lý và có thể hiểu được
                    4. Tránh:
                       - Copy nguyên văn từ context
                       - Sử dụng các biến đổi quá đơn giản
                       - Tạo câu không liên quan đến chủ đề
                       - Sử dụng thông tin ngoài context (trừ NEI có kiểm soát)
                    Trả về kết quả theo đúng format JSON với trường {self.claim_type}.
                    """

                # Sử dụng regex để xóa các khoảng trắng thừa
                prompt = re.sub(r'\s+', ' ', prompt).strip()
                # Xóa toàn bộ ký tự xuống dòng
                prompt = re.sub(r'\n', ' ', prompt).strip()
                response = model.generate_content(prompt)
                time.sleep(2)  # Prevent rate limiting
                return json.loads(response.text)["claims"]

            except Exception as e:
                error_msg = str(e)
                if "429" in error_msg or "500" in error_msg:
                    print(f"API error: {error_msg}")
                    self._rotate_api_key()
                    retries += 1
                else:
                    raise Exception(f"Unexpected error: {error_msg}")
        

    def generate_hard_claims(self, context: str, predict_func, max_attempts: int = 10) -> Dict[str, Any]:
        """Generate and validate claims with API key rotation support.
        Tracks and returns the claim with lowest probability across all attempts."""
        lowest_prob_claim = None
        lowest_prob = float('inf')  # Initialize with infinity to ensure first claim is captured
        attempts = 0
    
        while attempts < max_attempts:
            try:
                attempts += 1
                claims = self.generate_claims(context)
                
                for claim in claims:
                    prediction = predict_func(context, claim)
                    current_prob = prediction['probabilities'][self.claim_type]
                    other_labels = [label for label in ["SUPPORTED", "REFUTED", "NEI"] 
                                  if label != self.claim_type]
                    other_probs = [prediction['probabilities'][label] for label in other_labels]
                    
                    claim_result = {
                        'claim': claim,
                        'predicted_label': prediction['verdict'],
                        f'{self.claim_type}_probability': current_prob,
                        'other_probabilities': dict(zip(other_labels, other_probs)),
                        'all_probabilities': prediction['probabilities']
                    }
    
                    # Update lowest probability claim if current is lower
                    if current_prob < lowest_prob:
                        lowest_prob = current_prob
                        lowest_prob_claim = claim_result
    
                    # Validation logic based on label type
                    if self.claim_type == "NEI":
                        max_other_prob = max(other_probs)
                        if current_prob < sum(other_probs) or (0.7 > current_prob > 0.5):
                            return {
                                'success': True,
                                'attempts': attempts,
                                'results': [claim_result],
                                'used_fallback': False
                            }
                    else:
                        if current_prob < sum(other_probs):
                            return {
                                'success': True,
                                'attempts': attempts,
                                'results': [claim_result],
                                'used_fallback': False
                            }
    
            except Exception as e:
                print(f"Error during claim generation: {str(e)}")
                self._rotate_api_key()
                continue
    
        return {
            'success': True,
            'attempts': attempts,
            'results': [lowest_prob_claim] if lowest_prob_claim else [],
            'used_fallback': True
        }
        
def generate_claims_for_context(
    context: str,
    claim_types: List[str],
    predict_func,
    progress_bar
) -> List[Dict[str, Any]]:
    """Generate claims for the context with progress tracking and fallback handling."""
    claims_data = []
    total_steps = len(claim_types)
    
    for i, claim_type in enumerate(claim_types):
        progress_bar.progress((i / total_steps), f"Generating {claim_type} claims...")
        
        generator = ClaimGenerator(claim_type)
        result = generator.generate_hard_claims(context, predict_func)
        
        if result['success']:
            for claim_result in result['results']:
                claim_data = {
                    'context': context,
                    'claim': claim_result['claim'],
                    'label': claim_type,
                    'predicted_label': claim_result['predicted_label'],
                    f'{claim_type}_probability': claim_result[f'{claim_type}_probability'],
                    'probabilities': claim_result['all_probabilities']
                }
                claims_data.append(claim_data)
            
            if result.get('used_fallback'):
                st.warning(f"Used fallback claims for {claim_type} after {result['attempts']} attempts")
        else:
            st.error(f"Failed to generate any {claim_type} claims after {result['attempts']} attempts")
    
    return claims_data

def process_dataframe_with_claims(df: pd.DataFrame) -> pd.DataFrame:
    """Process each summary in the DataFrame and generate claims with fallback handling."""
    all_rows = []
    progress_bar = st.progress(0)
    
    for i, row in df.iterrows():
        context = row['Summary']
        
        for claim_type in ["SUPPORTED", "REFUTED", "NEI"]:
            generator = ClaimGenerator(claim_type)
            result = generator.generate_hard_claims(context, predict)
            
            if result['success']:
                for claim_result in result['results']:
                    new_row = row.copy()
                    new_row['claim'] = claim_result['claim']
                    new_row['label'] = claim_type
                    new_row['predicted_label'] = claim_result['predicted_label']
                    new_row['probability'] = claim_result[f'{claim_type}_probability']
                    new_row['all_probabilities'] = json.dumps(claim_result['all_probabilities'])
                    new_row['used_fallback'] = result.get('used_fallback', False)
                    all_rows.append(new_row)
                
                if result.get('used_fallback'):
                    st.warning(f"Used fallback claims for article {row['ID']} {claim_type}")
            
        progress_bar.progress((i + 1) / len(df))
    
    new_df = pd.DataFrame(all_rows)
    return new_df

def generate_claims_for_context(
    context: str,
    claim_types: List[str],
    predict_func,
    progress_bar
) -> List[Dict[str, Any]]:
    """Generate claims for the context with progress tracking."""
    claims_data = []
    total_steps = len(claim_types)
    
    for i, claim_type in enumerate(claim_types):
        progress_bar.progress((i / total_steps), f"Generating {claim_type} claims...")
        
        generator = ClaimGenerator(claim_type)
        result = generator.generate_hard_claims(context, predict_func)
        
        if result['success']:
            for claim_result in result['results']:
                claim_data = {
                    'context': context,
                    'claim': claim_result['claim'],
                    'label': claim_type,
                    'predicted_label': claim_result['predicted_label'],
                    f'{claim_type}_probability': claim_result[f'{claim_type}_probability'],
                    'probabilities': claim_result['all_probabilities']
                }
                claims_data.append(claim_data)
        else:
            st.error(f"Error generating {claim_type} claims after {result['attempts']} attempts")
    
    return claims_data

# Wikipedia helper functions
def format_content(content):
    cleaned_content = re.sub(r'==.*?==', '', content)
    cleaned_text = re.sub(r'\n\s*\n', '\n', cleaned_content)
    lines = cleaned_text.split("\n")
    cleaned_lines = [line.strip() for line in lines]
    return "\n\n".join(cleaned_lines)

def chars_to_nums(s):
    return ''.join(str((char_to_num.get(char.upper(), 0) % 4) + 1) for char in s if char.upper() in char_to_num)

def create_id(topic, stt):
    topic_abbr = ''.join([word[0] for word in topic.split() if word])
    topic_num = chars_to_nums(topic_abbr)
    return f'uit_{topic_num}_{stt}'

def wikipedia_scrape(title_input, stt, filename):
    try:
        page = wikipedia.page(title_input)
        title = page.title
        summary = wikipedia.summary(title_input, sentences=10)
        url = page.url
        topic = filename.split(".")[0]
        base_article_id = create_id(topic, stt)

        summary = format_content(summary)
        # Chia câu dựa trên dấu chấm
        sentences = split_sentences_vietnamese(summary)
        total_sentences = len(sentences)

        # Tính toán chia đoạn động
        part_lengths = [5, total_sentences - 5]
        
        results = []
        start = 0
        for i, length in enumerate(part_lengths, 1):
            # Đảm bảo không vượt quá tổng số câu
            end = min(start + length, total_sentences)
            part = '. '.join(sentences[start:end]) + '.'

            if len(part)>=5:
                article_id = f"{base_article_id}_{i}"
                results.append({
                    "ID": article_id,
                    "Title": title,
                    "Topic": topic,
                    "Summary": part,
                    "URL": url
                })
                
                start = end

        return results
    except (wikipedia.exceptions.DisambiguationError, wikipedia.exceptions.PageError) as e:
        return None



@st.cache_data
def convert_df_to_csv(df):
    return df.to_csv(index=False).encode("utf-8-sig")

def process_dataframe_with_claims(df: pd.DataFrame) -> pd.DataFrame:
    """Process each summary in the DataFrame and generate claims."""
    
    # Create lists to store the new data
    all_rows = []
    progress_bar = st.progress(0)
    # Process each row in the original DataFrame
    for i, row in df.iterrows():
        context = row['Summary']
        
        # Generate claims for each type
        for claim_type in ["SUPPORTED", "REFUTED", "NEI"]:
            generator = ClaimGenerator(claim_type)
            result = generator.generate_hard_claims(context, predict)
            
            if result['success']:
                for claim_result in result['results']:
                    # Create a new row with all original columns plus new claim data
                    new_row = row.copy()
                    new_row['claim'] = claim_result['claim']
                    new_row['label'] = claim_type
                    new_row['predicted_label'] = claim_result['predicted_label']
                    new_row['probability'] = claim_result[f'{claim_type}_probability']
                    new_row['support_probability'] = claim_result['all_probabilities']["SUPPORTED"]
                    new_row['refuted_probability'] = claim_result['all_probabilities']["REFUTED"]
                    new_row['nei_probability'] = claim_result['all_probabilities']["NEI"]
                    all_rows.append(new_row)
            progress_bar.progress((i + 1) / len(df))
            
    
    # Create new DataFrame with all rows
    new_df = pd.DataFrame(all_rows)
    return new_df

def main():
    st.title("Ứng dụng crawl data wiki tiếng việt và tạo claims")

    uploaded_file = st.file_uploader("Tải lên tệp văn bản (txt)", type=["txt"])

    if uploaded_file is not None:
        file_progress = st.progress(0)
        status_text = st.empty()
        
        file_contents = uploaded_file.getvalue().decode("utf-8")
        titles = file_contents.split("\n")
        
        articles_info = []
        for index, title in enumerate(titles):
            title = title.strip()
            if title:
                status_text.text(f"Processing article {index + 1}/{len(titles)}: {title}")
                article_infos = wikipedia_scrape(title, index + 1, uploaded_file.name)
                
                # Safely append results, handling potential None values
                if article_infos:
                    articles_info.extend(article_infos)
                file_progress.progress((index + 1) / len(titles))
            else:
                st.warning(f"Dòng {index + 1}: Tiêu đề trống")
        
        file_progress.empty()
        status_text.empty()

        if not articles_info:
            st.error("Không có bài viết nào được thu thập từ Wikipedia.")
        else:
            df = pd.DataFrame(articles_info)
            
            st.subheader("Danh sách bài viết từ Wikipedia:")
            st.dataframe(df, use_container_width=True, hide_index=True)

            # Add automatic claim generation for all summaries
            if st.button("Tạo Claims cho tất cả bài viết"):
                with st.spinner("Đang tạo claims cho tất cả bài viết..."):
                    
                    # Process the DataFrame and generate claims
                    df_with_claims = process_dataframe_with_claims(df)
                    
                    st.subheader("Bảng dữ liệu với Claims:")
                    st.dataframe(df_with_claims, use_container_width=True)
                    
                    # Download options
                    csv = convert_df_to_csv(df_with_claims)
                    claims_filename = f"uit_{uploaded_file.name.split('.')[0]}_with_claims.csv"
                    st.download_button(
                        label="Download CSV với Claims",
                        data=csv,
                        file_name=claims_filename,
                        mime="text/csv",
                    )
            
            # Original download option for basic DataFrame
            csv = convert_df_to_csv(df)
            csv_filename = "uit_" + uploaded_file.name.split(".")[0] + ".csv"
            st.download_button(
                label="Download CSV (không có claims)",
                data=csv,
                file_name=csv_filename,
                mime="text/csv",
            )

if __name__ == "__main__":
    main()