src/demo/main.py

import csv
import json
import os

import pandas as pd
from langchain_community.embeddings import HuggingFaceEmbeddings
from asg_retriever import legal_pdf
from asg_loader import DocumentLoading
from asg_retriever import Retriever, query_embeddings_new_new
from asg_generator import generate_sentence_patterns, generate
from category_and_tsne import clustering
from langchain_text_splitters  import RecursiveCharacterTextSplitter
import time
import torch
import re
import transformers
from dotenv import load_dotenv
from asg_clustername import generate_cluster_name_new
from asg_outline import OutlineGenerator, generateSurvey_qwen_new
import os
from markdown_pdf import MarkdownPdf, Section  # Assuming you are using markdown_pdf
from typing import Any
from .path_utils import get_path, setup_hf_cache

# 设置 Hugging Face 缓存目录
cache_dir = setup_hf_cache()

def clean_str(input_str):
    input_str = str(input_str).strip().lower()
    if input_str == "none" or input_str == "nan" or len(input_str) == 0:
        return ""
    input_str = input_str.replace('\\n',' ').replace('\n',' ').replace('\r',' ').replace('——',' ').replace('——',' ').replace('__',' ').replace('__',' ').replace('........','.').replace('....','.').replace('....','.').replace('..','.').replace('..','.').replace('..','.').replace('. . . . . . . . ','. ').replace('. . . . ','. ').replace('. . . . ','. ').replace('. . ','. ').replace('. . ','. ')
    input_str = re.sub(r'\\u[0-9a-z]{4}', ' ', input_str).replace('  ',' ').replace('  ',' ')
    return input_str

def remove_invalid_citations(text, valid_collection_names):
    """
    只保留 [xxx\] 中的 xxx 属于 valid_collection_names 的引用，
    其余的引用标记一律删除。
    """
    pattern = r"\[(.*?)\\\]"  # 匹配形如 [xxx\] 的内容
    all_matches = re.findall(pattern, text)

    new_text = text
    for match in all_matches:
        cleaned_match = match.rstrip('\\')  # 去除末尾的 \
        if cleaned_match not in valid_collection_names:
            new_text = new_text.replace(f"[{match}\\]", "")
    return new_text
def normalize_citations_with_mapping(paper_text):
    # 使用正则表达式匹配所有引用标记（形如 [citation1]）
    citations = re.findall(r'\[.*?\]', paper_text)
    # 去重并保持顺序
    unique_citations = list(dict.fromkeys(citations))
    # 生成引用映射表，把原始引用标记映射为数字引用
    citation_mapping = {citation: f'[{i + 1}]' for i, citation in enumerate(unique_citations)}

    # 在文本中替换老引用为新引用
    normalized_text = paper_text
    for old_citation, new_citation in citation_mapping.items():
        normalized_text = normalized_text.replace(old_citation, new_citation)

    # 生成从数字到原始引用标记的反向映射
    # 用 rstrip('\\') 去掉末尾的反斜杠
    reverse_mapping = {
        i + 1: unique_citations[i].strip('[]').rstrip('\\')
        for i in range(len(unique_citations))
    }

    return normalized_text, reverse_mapping
def generate_references_section(citation_mapping, collection_pdf_mapping):
    
    references = ["# References"]  # 生成引用部分
    for num in sorted(citation_mapping.keys()):
        collection_name = citation_mapping[num]
        pdf_name = collection_pdf_mapping.get(collection_name, "Unknown PDF")
        if pdf_name.endswith(".pdf"):
            pdf_name = pdf_name[:-4]
        # 在每一行末尾添加两个空格以确保换行
        references.append(f"[{num}] {pdf_name}  ")

    return "\n".join(references)
def fix_citation_punctuation_md(text):
    """
    把类似于 'some text. \[1]' 或 'some text. \[2]' 调整为 'some text \[1].'
    仅针对已经变成 \[1], \[2] 之类数字引用的 Markdown 情况有效。
    如果还没有变成 \[数字]，则需先经过 normalize_citations_with_mapping。
    """
    # 正则表达式匹配点号后带有空格或无空格，紧接 \[数字] 的情况
    pattern = r'\.\s*(\\\[\d+\])'
    replacement = r' \1.'
    fixed_text = re.sub(pattern, replacement, text)
    return fixed_text
def finalize_survey_paper(paper_text, 
                          Global_collection_names, 
                          Global_file_names):

    # 1) 删除所有不想要的旧引用（包括 [数字]、[Sewon, 2021] 等）
    paper_text = remove_invalid_citations(paper_text, Global_collection_names)

    # 2) 规范化引用 => [1][2]...
    normalized_text, citation_mapping = normalize_citations_with_mapping(paper_text)
    
    # 3) 修复标点，比如 .[1] => [1].
    normalized_text = fix_citation_punctuation_md(normalized_text)

    # 4) 构造 {collection_name: pdf_file_name} 字典
    collection_pdf_mapping = dict(zip(Global_collection_names, Global_file_names))
    
    # 5) 生成 References
    references_section = generate_references_section(citation_mapping, collection_pdf_mapping)
    
    # 6) 合并正文和 References
    final_paper = normalized_text.strip() + "\n\n" + references_section
    return final_paper

class ASG_system:
    def __init__(self, root_path: str, survey_id:str, pdf_path: str, survey_title: str, cluster_standard: str) -> None:
        load_dotenv()
        self.pdf_path = pdf_path
        self.txt_path = root_path + "/txt"
        self.tsv_path = root_path + "/tsv"
        self.md_path = root_path + "/md"
        self.info_path = root_path + "/info"
        self.result_path = root_path + "/result"

        self.survey_id = survey_id
        self.survey_title = survey_title
        self.cluster_standard = cluster_standard

        self.collection_names = []
        self.file_names = []
        self.citation_data = []
        self.description_list = []
        self.ref_list = []
        self.cluster_names = []
        self.collection_names_clustered = []
        self.df_selected = ''


        model_id = "meta-llama/Meta-Llama-3.1-8B-Instruct"
        try:
            self.embedder = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2", cache_folder=cache_dir)
        except Exception as e:
            print(f"Error initializing embedder: {e}")
            self.embedder = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
        self.pipeline = transformers.pipeline(
            "text-generation",
            model=model_id,
            model_kwargs={"torch_dtype": torch.bfloat16},
            token = os.getenv('HF_API_KEY'),
            device_map="auto",
        )
        self.pipeline.model.load_adapter(peft_model_id = "technicolor/llama3.1_8b_outline_generation", adapter_name="outline")
        self.pipeline.model.load_adapter(peft_model_id ="technicolor/llama3.1_8b_abstract_generation", adapter_name="abstract")
        self.pipeline.model.load_adapter(peft_model_id ="technicolor/llama3.1_8b_conclusion_generation", adapter_name="conclusion")

        os.makedirs(self.txt_path, exist_ok=True)
        os.makedirs(f'{self.txt_path}/{self.survey_id}', exist_ok=True)

        os.makedirs(self.tsv_path, exist_ok=True)

        os.makedirs(self.md_path, exist_ok=True)
        os.makedirs(f'{self.md_path}/{self.survey_id}', exist_ok=True)

        os.makedirs(self.info_path, exist_ok=True)
        os.makedirs(f'{self.info_path}/{self.survey_id}', exist_ok=True)

        os.makedirs(self.result_path, exist_ok=True)
        os.makedirs(f'{self.result_path}/{self.survey_id}', exist_ok=True)

    def parsing_pdfs(self, mode="intro") -> None:
        pdf_files = os.listdir(self.pdf_path)
        loader = DocumentLoading()
        

        for pdf_file in pdf_files:

            pdf_file = os.path.join(self.pdf_path, pdf_file)

            split_start_time = time.time()

            base_name = os.path.splitext(os.path.basename(pdf_file))[0]
            target_dir = os.path.join(self.md_path, self.survey_id, base_name, "auto")
            md_dir = os.path.join(self.md_path, self.survey_id)

            loader.convert_pdf_to_md(pdf_file, md_dir)

            md_file_path = os.path.join(target_dir, f"{base_name}.md")
            print(md_file_path)
            print("*"*24)
            if not os.path.exists(md_file_path):
                raise FileNotFoundError(f"Markdown file {md_file_path} does not exist. Conversion might have failed.")

            if mode == "intro":
                doc = loader.process_md_file(md_file_path, self.survey_id, self.txt_path)
            elif mode == "full":
                doc = loader.process_md_file_full(md_file_path, self.survey_id,self.txt_path)
            
            text_splitter = RecursiveCharacterTextSplitter(
                chunk_size=400,
                chunk_overlap=30,
                length_function=len,
                is_separator_regex=False,
            )
            splitters = text_splitter.create_documents([doc])
            documents_list = [document.page_content for document in splitters]
            for i in range(len(documents_list)):
                documents_list[i] = documents_list[i].replace('\n', ' ')
            print(f"Splitting took {time.time() - split_start_time} seconds.")

            embed_start_time = time.time()

            doc_results = self.embedder.embed_documents(documents_list)
            if isinstance(doc_results, torch.Tensor):
                embeddings_list = doc_results.tolist()
            else:
                embeddings_list = doc_results
            print(f"Embedding took {time.time() - embed_start_time} seconds.")

            # Prepare metadata
            metadata_list = [{"doc_name": os.path.basename(pdf_file)} for i in range(len(documents_list))]

            title = os.path.splitext(os.path.basename(pdf_file))[0]
            

            title_new = title.strip()
            invalid_chars = ['<', '>', ':', '"', '/', '\\', '|', '?', '*','_']
            for char in invalid_chars:
                title_new = title_new.replace(char, ' ')
            print("============================")
            print(title_new)

            # New logic to create collection_name
            # filename = os.path.basename(file_path)
            collection_name = legal_pdf(title_new)

            retriever = Retriever()
            retriever.list_collections_chroma()
            retriever.create_collection_chroma(collection_name)
            retriever.add_documents_chroma(
                collection_name=collection_name,
                embeddings_list=embeddings_list,
                documents_list=documents_list,
                metadata_list=metadata_list
            )

            self.collection_names.append(collection_name)
            self.file_names.append(title_new)
        print(self.collection_names)
        print(self.file_names)

        json_files = os.listdir(os.path.join(self.txt_path, self.survey_id))
        ref_paper_num = len(json_files)
        print(f'The length of the json files is {ref_paper_num}')


        json_data_pd = pd.DataFrame()   
        for _ in json_files:
            file_path = os.path.join(self.txt_path, self.survey_id, _)

            with open(file_path, 'r', encoding="utf-8") as file:
                data = json.load(file)

                # Extract necessary information
                title = data.get("title", "")
                abstract = data.get("abstract", "")
                authors = data.get("authors", "")
                introduction = data.get("introduction", "")
                new_data = {
                    "reference paper title": title,
                    "reference paper citation information (can be collected from Google scholar/DBLP)": authors,
                    "reference paper abstract (Please copy the text AND paste here)": abstract,
                    "reference paper introduction (Please copy the text AND paste here)": introduction,
                    "reference paper doi link (optional)": "",
                    "reference paper category label (optional)": ""
                }

                # 将新数据转换为 DataFrame
                new_data_df = pd.DataFrame([new_data])

                # 使用 pd.concat 而不是 append
                json_data_pd = pd.concat([json_data_pd, new_data_df], ignore_index=True)

        # Save the DataFrame to a variable for further use
        input_pd = json_data_pd

        if ref_paper_num>0:
                
            ## change col name
            input_pd['ref_title'] = [filename for filename in self.file_names]
            input_pd["ref_context"] = [""]*ref_paper_num
            input_pd["ref_entry"] = input_pd["reference paper citation information (can be collected from Google scholar/DBLP)"]
            input_pd["abstract"] = input_pd["reference paper abstract (Please copy the text AND paste here)"].apply(lambda x: clean_str(x) if len(str(x))>0 else 'Invalid abstract')
            input_pd["intro"] = input_pd["reference paper introduction (Please copy the text AND paste here)"].apply(lambda x: clean_str(x) if len(str(x))>0 else 'Invalid introduction')

            # optional columns
            input_pd["label"] = input_pd["reference paper category label (optional)"].apply(lambda x: str(x) if len(str(x))>0 else '')
            #input_pd["label"] = input_pd["reference paper category id (optional)"].apply(lambda x: str(x) if len(str(x))>0 else '')
            ## output tsv
            # output_tsv_filename = self.tsv_path + self.survey_id + '.tsv'
            output_tsv_filename = os.path.join(self.tsv_path, self.survey_id + '.tsv') 

            #output_df = input_pd[["ref_title","ref_context","ref_entry","abstract","intro","description"]]
            output_df = input_pd[["ref_title","ref_context","ref_entry","abstract","intro", 'label']]
            # print(output_df)

            #pdb.set_trace()
            output_df.to_csv(output_tsv_filename, sep='\t')

    def description_generation(self) -> None:
        query= self.cluster_standard
        query_list = generate_sentence_patterns(query)
        for name in self.collection_names:
            context, citation_data = query_embeddings_new_new(name, query_list)
            self.citation_data.extend(citation_data)

            description = generate(context, query, name)
            self.description_list.append(description)

        citation_path = f'{self.info_path}/{self.survey_id}/citation_data.json'
        os.makedirs(os.path.dirname(citation_path), exist_ok=True)
        with open(citation_path, 'w', encoding="utf-8") as outfile:
            json.dump(self.citation_data, outfile, indent=4, ensure_ascii=False)
        
        file_path = f'{self.tsv_path}/{self.survey_id}.tsv'

        with open(file_path, 'r', newline='', encoding='utf-8') as infile:
            reader = csv.reader(infile, delimiter='\t')
            rows = list(reader)
        if rows:
            headers = rows[0]
            headers.append('retrieval_result')

            updated_rows = [headers]
            for row, description in zip(rows[1:], self.description_list):
                row.append(description)
                updated_rows.append(row)

            with open(file_path, 'w', newline='', encoding='utf-8') as outfile:
                writer = csv.writer(outfile, delimiter='\t')
                writer.writerows(updated_rows)

            print('Updated file has been saved to', file_path)
        else:
            print('Input file is empty.')

    def agglomerative_clustering(self) -> None:
        df = pd.read_csv(f'{self.tsv_path}/{self.survey_id}.tsv', sep='\t', index_col=0, encoding='utf-8')
        df_selected = df

        df_selected, _ = clustering(df_selected, 3, self.survey_id, self.info_path, self.tsv_path)
        self.df_selected = df_selected

        df_tmp = df_selected.reset_index()
        df_tmp['index'] = df_tmp.index
        ref_titles = list(df_tmp.groupby(df_tmp['label'])['ref_title'].apply(list))
        # ref_indexs = list(df_tmp.groupby(df_tmp['label'])['index'].apply(list))

        category_label_summarized = generate_cluster_name_new(f"{self.tsv_path}/{self.survey_id}.tsv", self.survey_title)
        self.cluster_names = category_label_summarized

        cluster_info = {category_label_summarized[i]:ref_titles[i] for i in range(len(category_label_summarized))}
        for key, value in cluster_info.items():
            temp = [legal_pdf(i) for i in value]
            cluster_info[key] = temp
            self.collection_names_clustered.append(temp)
        cluster_info_path = f'{self.info_path}/{self.survey_id}/cluster_info.json'
        with open(cluster_info_path, 'w', encoding="utf-8") as outfile:
            json.dump(cluster_info, outfile, indent=4, ensure_ascii=False)        

    def outline_generation(self) -> None:
        print(self.df_selected)
        print(self.cluster_names)
        outline_generator = OutlineGenerator(self.pipeline, self.df_selected, self.cluster_names)
        outline_generator.get_cluster_info()
        messages, outline = outline_generator.generate_outline_qwen(self.survey_title)
        outline_json = {'messages':messages, 'outline': outline}
        output_path = f'{self.info_path}/{self.survey_id}/outline.json'
        os.makedirs(os.path.dirname(output_path), exist_ok=True)
        with open(output_path, 'w', encoding="utf-8") as outfile:
            json.dump(outline_json, outfile, indent=4, ensure_ascii=False)
        

    def section_generation(self) -> None:
        generateSurvey_qwen_new(self.survey_id, self.survey_title, self.collection_names_clustered, self.pipeline, self.citation_data, './txt','./info')

    def citation_generation(self) -> None:
        """
        Generate citation Markdown and PDF files from JSON and store them in the specified result path.
        """

        json_filepath = os.path.join(self.info_path, self.survey_id, "generated_result.json")

        markdown_dir = f'{self.result_path}/{self.survey_id}'
        markdown_filename = f'survey_{self.survey_id}.md'
        markdown_filepath = os.path.join(markdown_dir, markdown_filename)
        pdf_filename = f'survey_{self.survey_id}.pdf'
        pdf_filepath = os.path.join(markdown_dir, pdf_filename)

        markdown_content = self.get_markdown_content(json_filepath)
        if not markdown_content:
            raise ValueError("Markdown content is empty. Cannot generate citation files.")

        try:
            with open(markdown_filepath, 'w', encoding='utf-8', encoding="utf-8") as markdown_file:
                markdown_file.write(markdown_content)
            print(f"Markdown content saved to: {markdown_filepath}")
        except Exception as e:
            raise RuntimeError(f"Failed to save Markdown file: {e}")

        try:
            pdf = MarkdownPdf()
            pdf.meta["title"] = "Citation Results" 
            pdf.add_section(Section(markdown_content, toc=False))  
            pdf.save(pdf_filepath) 
            print(f"PDF content saved to: {pdf_filepath}")
        except Exception as e:
            raise RuntimeError(f"Failed to generate PDF file: {e}")
        print(f"Files generated successfully: \nMarkdown: {markdown_filepath}\nPDF: {pdf_filepath}")

    def get_markdown_content(self, json_filepath: str) -> str:
        """
        Read a JSON file and generate Markdown content based on its data.

        :param json_filepath: Path to the JSON file containing survey data.
        :return: A string containing the generated Markdown content.
        """
        try:
            with open(json_filepath, 'r', encoding='utf-8', encoding="utf-8") as json_file:
                survey_data = json.load(json_file)
        except Exception as e:
            raise RuntimeError(f"Failed to read JSON file: {e}")

        topic = survey_data.get('survey_title', 'Default Topic')
        content = survey_data.get('content', 'No content available.')

        survey_title_markdown = f"# A Survey of {topic}\n\n"
        survey_content_markdown = content + "\n\n"

        markdown_content = survey_title_markdown + survey_content_markdown
        markdown_content = finalize_survey_paper(markdown_content, self.collection_names, self.file_names)
        return markdown_content

if __name__ == "__main__":
    root_path = "."
    pdf_path = "./pdfs/test"
    survey_title = "Automating Literature Review Generation with LLM"
    cluster_standard = "method"
    asg_system = ASG_system(root_path, 'test', pdf_path, survey_title, cluster_standard)
    asg_system.parsing_pdfs()
    asg_system.description_generation()
    asg_system.agglomerative_clustering()
    asg_system.outline_generation()
    asg_system.section_generation()
    asg_system.citation_generation()