update cache

src/demo/asg_retriever.py

@@ -8,7 +8,11 @@ from .asg_splitter import TextSplitting
 from langchain_huggingface import HuggingFaceEmbeddings
 import time
 import concurrent.futures
-from .path_utils import get_path
+from .path_utils import get_path, setup_hf_cache
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+
+# 设置 Hugging Face 缓存目录
+cache_dir = setup_hf_cache()
 
 class Retriever:
     client = None
@@ -201,7 +205,11 @@ def process_pdf(file_path: str, survey_id: str, embedder: HuggingFaceEmbeddings,
     return collection_name, embeddings_list, documents_list, metadata_list,title_new
 
 def query_embeddings(collection_name: str, query_list: list):
-    embedder = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
+    try:
+        embedder = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2", cache_folder=cache_dir)
+    except Exception as e:
+        print(f"Error initializing embedder: {e}")
+        embedder = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
     retriever = Retriever()
 
     final_context = ""
@@ -222,7 +230,11 @@ def query_embeddings(collection_name: str, query_list: list):
 
 # new, may be in parallel
 def query_embeddings_new(collection_name: str, query_list: list):
-    embedder = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
+    try:
+        embedder = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2", cache_folder=cache_dir)
+    except Exception as e:
+        print(f"Error initializing embedder: {e}")
+        embedder = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
     retriever = Retriever()
 
     final_context = ""
@@ -250,7 +262,11 @@ def query_embeddings_new(collection_name: str, query_list: list):
 
 # wza
 def query_embeddings_new_new(collection_name: str, query_list: list):
-    embedder = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
+    try:
+        embedder = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2", cache_folder=cache_dir)
+    except Exception as e:
+        print(f"Error initializing embedder: {e}")
+        embedder = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
     retriever = Retriever()
 
     final_context = ""  # Stores concatenated context
@@ -313,7 +329,11 @@ def query_multiple_collections(collection_names: list[str], query_list: list[str
         dict: Combined results from all collections, grouped by collection.
     """
     # Define embedder inside the function
-    embedder = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
+    try:
+        embedder = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2", cache_folder=cache_dir)
+    except Exception as e:
+        print(f"Error initializing embedder: {e}")
+        embedder = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
     retriever = Retriever()
 
     def query_single_collection(collection_name: str):

src/demo/category_and_tsne.py

@@ -7,6 +7,8 @@ import seaborn as sns
 import json
 from sklearn.manifold import TSNE
 from sklearn.cluster import AgglomerativeClustering
+import os
+import tempfile
 
 from sentence_transformers import SentenceTransformer
 from bertopic import BERTopic
@@ -14,7 +16,7 @@ from bertopic.representation import KeyBERTInspired
 from sklearn.feature_extraction.text import CountVectorizer
 from bertopic.vectorizers import ClassTfidfTransformer
 from umap import UMAP
-from .path_utils import get_path
+from .path_utils import get_path, setup_hf_cache
 
 plt.switch_backend('agg')
 device = 0
@@ -35,6 +37,9 @@ import matplotlib.pyplot as plt
 from sklearn.manifold import TSNE
 import seaborn as sns
 
+# 设置 Hugging Face 缓存目录
+cache_dir = setup_hf_cache()
+
 class DimensionalityReduction:
     def fit(self, X):
         return self
@@ -44,7 +49,11 @@ class DimensionalityReduction:
 
 class ClusteringWithTopic:
     def __init__(self, df, n_topics=3):
-        embedding_model = SentenceTransformer("nomic-ai/nomic-embed-text-v1", trust_remote_code=True)
+        try:
+            embedding_model = SentenceTransformer("nomic-ai/nomic-embed-text-v1", trust_remote_code=True, cache_folder=cache_dir)
+        except Exception as e:
+            print(f"Error initializing SentenceTransformer: {e}")
+            embedding_model = SentenceTransformer("nomic-ai/nomic-embed-text-v1", trust_remote_code=True)
         # umap_model = DimensionalityReduction()
         umap_model = UMAP(n_neighbors=15, n_components=5, min_dist=0.0, metric='cosine', init = 'pca')
         hdbscan_model = AgglomerativeClustering(n_clusters=n_topics)
@@ -81,7 +90,11 @@ class ClusteringWithTopic:
         初始化 ClusteringWithTopic，接受一个 n_topics_list，其中包含多个聚类数目，
         选取 silhouette_score 最高的结果。
         """
-        embedding_model = SentenceTransformer("nomic-ai/nomic-embed-text-v1", trust_remote_code=True)
+        try:
+            embedding_model = SentenceTransformer("nomic-ai/nomic-embed-text-v1", trust_remote_code=True, cache_folder=cache_dir)
+        except Exception as e:
+            print(f"Error initializing SentenceTransformer: {e}")
+            embedding_model = SentenceTransformer("nomic-ai/nomic-embed-text-v1", trust_remote_code=True)
         self.embeddings = embedding_model.encode(df, show_progress_bar=True)
         
         self.df = df
@@ -97,7 +110,7 @@ class ClusteringWithTopic:
         # 用于存储不同聚类数目的结果
         self.best_n_topics = None
         self.best_labels = None
-        self.best_score = -1 
+        self.best_score = -1
     # def fit_and_get_labels(self, X):
     #     topics, probs = self.topic_model.fit_transform(self.df, self.embeddings)
     #     return topics 

src/demo/main.py

@@ -20,6 +20,10 @@ 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()
@@ -135,7 +139,11 @@ class ASG_system:
 
 
         model_id = "meta-llama/Meta-Llama-3.1-8B-Instruct"
-        self.embedder = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
+        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,

src/demo/path_utils.py

@@ -1,6 +1,21 @@
 import os
 import tempfile
 
+# 设置 Hugging Face 缓存目录
+def setup_hf_cache():
+    """设置 Hugging Face 缓存目录，在 Hugging Face Spaces 中使用临时目录"""
+    if os.environ.get('SPACE_ID') or os.environ.get('HF_SPACE_ID'):
+        # 在 Hugging Face Spaces 中使用临时目录作为缓存
+        cache_dir = tempfile.mkdtemp()
+        os.environ['HF_HOME'] = cache_dir
+        os.environ['TRANSFORMERS_CACHE'] = os.path.join(cache_dir, 'transformers')
+        os.environ['HF_HUB_CACHE'] = os.path.join(cache_dir, 'hub')
+        print(f"Using Hugging Face cache directory: {cache_dir}")
+        return cache_dir
+    else:
+        # 本地环境使用默认缓存目录
+        return None
+
 # 检查是否在 Hugging Face Spaces 环境中
 def get_data_paths():
     # 如果在 Hugging Face Spaces 中，使用临时目录

src/demo/survey_generation_pipeline/asg_retriever.py

@@ -8,7 +8,11 @@ from .asg_splitter import TextSplitting
 from langchain_huggingface import HuggingFaceEmbeddings
 import time
 import concurrent.futures
-from ..path_utils import get_path
+from ..path_utils import get_path, setup_hf_cache
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+
+# 设置 Hugging Face 缓存目录
+cache_dir = setup_hf_cache()
 
 class Retriever:
     client = None
@@ -223,7 +227,11 @@ def process_pdf(file_path: str, survey_id: str, embedder: HuggingFaceEmbeddings,
     return collection_name, embeddings_list, documents_list, metadata_list,title_new
 
 def query_embeddings(collection_name: str, query_list: list):
-    embedder = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
+    try:
+        embedder = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2", cache_folder=cache_dir)
+    except Exception as e:
+        print(f"Error initializing embedder: {e}")
+        embedder = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
     retriever = Retriever()
 
     final_context = ""
@@ -244,7 +252,11 @@ def query_embeddings(collection_name: str, query_list: list):
 
 # new, may be in parallel
 def query_embeddings_new(collection_name: str, query_list: list):
-    embedder = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
+    try:
+        embedder = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2", cache_folder=cache_dir)
+    except Exception as e:
+        print(f"Error initializing embedder: {e}")
+        embedder = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
     retriever = Retriever()
 
     final_context = ""
@@ -270,45 +282,59 @@ def query_embeddings_new(collection_name: str, query_list: list):
                     seen_chunks.add(chunk)
     return final_context
 
-def query_embeddings_new_new(collection_name: str, query_list: list, retriever):
-    embedder = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
-    final_context = ""
-    citation_data_list = []
-    seen_chunks = set()
+# wza
+def query_embeddings_new_new(collection_name: str, query_list: list):
+    try:
+        embedder = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2", cache_folder=cache_dir)
+    except Exception as e:
+        print(f"Error initializing embedder: {e}")
+        embedder = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
+    retriever = Retriever()
 
-    for query_text in query_list:
-        try:
-            query_embeddings = embedder.embed_query(query_text)
-            query_result = retriever.query_chroma(
-                collection_name=collection_name,
-                query_embeddings=[query_embeddings],
-                n_results=5
-            )
-        except Exception as e:
-            print(f"Query '{query_text}' failed with exception: {e}")
-            continue
-
-        if "documents" not in query_result or "distances" not in query_result:
-            continue
-        if not query_result["documents"] or not query_result["distances"]:
-            continue
-
-        docs_list = query_result["documents"][0] if query_result["documents"] else []
-        dist_list = query_result["distances"][0] if query_result["distances"] else []
-
-        if len(docs_list) != len(dist_list):
-            continue
-
-        for chunk, distance in zip(docs_list, dist_list):
-            processed_chunk = chunk.strip()
-            if processed_chunk not in seen_chunks:
-                final_context += processed_chunk + "//\n"
-                seen_chunks.add(processed_chunk)
-                citation_data_list.append({
-                    "source": collection_name,
-                    "distance": distance,
-                    "content": processed_chunk,
-                })
+    final_context = ""  # Stores concatenated context
+    citation_data_list = []  # Stores chunk content and collection name as source
+    seen_chunks = set()  # Ensures unique chunks are added
+
+    def process_query(query_text):
+        # Embed the query text and retrieve relevant chunks
+        query_embeddings = embedder.embed_query(query_text)
+        query_result = retriever.query_chroma(
+            collection_name=collection_name,
+            query_embeddings=[query_embeddings],
+            n_results=5  # Fixed number of results
+        )
+        return query_result
+
+    with concurrent.futures.ThreadPoolExecutor() as executor:
+        future_to_query = {executor.submit(process_query, q): q for q in query_list}
+        for future in concurrent.futures.as_completed(future_to_query):
+            query_text = future_to_query[future]
+            try:
+                query_result = future.result()
+            except Exception as e:
+                print(f"Query '{query_text}' failed with exception: {e}")
+                continue
+
+            if "documents" not in query_result or "distances" not in query_result:
+                continue
+            if not query_result["documents"] or not query_result["distances"]:
+                continue
+            docs_list = query_result["documents"][0] if query_result["documents"] else []
+            dist_list = query_result["distances"][0] if query_result["distances"] else []
+
+            if len(docs_list) != len(dist_list):
+                continue
+
+            for chunk, distance in zip(docs_list, dist_list):
+                processed_chunk = chunk.strip()
+                if processed_chunk not in seen_chunks:
+                    final_context += processed_chunk + "//\n"
+                    seen_chunks.add(processed_chunk)
+                    citation_data_list.append({
+                        "source": collection_name,
+                        "distance": distance,
+                        "content": processed_chunk,
+                    })
 
     return final_context, citation_data_list
 
@@ -325,7 +351,11 @@ def query_multiple_collections(collection_names: list[str], query_list: list[str
         dict: Combined results from all collections, grouped by collection.
     """
     # Define embedder inside the function
-    embedder = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
+    try:
+        embedder = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2", cache_folder=cache_dir)
+    except Exception as e:
+        print(f"Error initializing embedder: {e}")
+        embedder = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
     retriever = Retriever()
 
     def query_single_collection(collection_name: str):

src/demo/survey_generation_pipeline/category_and_tsne.py

@@ -1,15 +1,22 @@
 from sklearn.metrics import silhouette_score
 
 import numpy as np
+import pandas as pd
 import matplotlib.pyplot as plt
 import seaborn as sns
-import matplotlib.pyplot as plt
+import json
 from sklearn.manifold import TSNE
 from sklearn.cluster import AgglomerativeClustering
-import json
-from ..path_utils import get_path
+import os
+import tempfile
 
-IMG_PATH = get_path('img')
+from sentence_transformers import SentenceTransformer
+from bertopic import BERTopic
+from bertopic.representation import KeyBERTInspired
+from sklearn.feature_extraction.text import CountVectorizer
+from bertopic.vectorizers import ClassTfidfTransformer
+from umap import UMAP
+from ..path_utils import get_path, setup_hf_cache
 
 plt.switch_backend('agg')
 device = 0
@@ -30,6 +37,9 @@ import matplotlib.pyplot as plt
 from sklearn.manifold import TSNE
 import seaborn as sns
 
+# 设置 Hugging Face 缓存目录
+cache_dir = setup_hf_cache()
+
 class DimensionalityReduction:
     def fit(self, X):
         return self
@@ -39,7 +49,11 @@ class DimensionalityReduction:
 
 class ClusteringWithTopic:
     def __init__(self, df, n_topics=3):
-        embedding_model = SentenceTransformer("nomic-ai/nomic-embed-text-v1", trust_remote_code=True)
+        try:
+            embedding_model = SentenceTransformer("nomic-ai/nomic-embed-text-v1", trust_remote_code=True, cache_folder=cache_dir)
+        except Exception as e:
+            print(f"Error initializing SentenceTransformer: {e}")
+            embedding_model = SentenceTransformer("nomic-ai/nomic-embed-text-v1", trust_remote_code=True)
         # umap_model = DimensionalityReduction()
         umap_model = UMAP(n_neighbors=15, n_components=5, min_dist=0.0, metric='cosine', init = 'pca')
         hdbscan_model = AgglomerativeClustering(n_clusters=n_topics)
@@ -76,7 +90,11 @@ class ClusteringWithTopic:
         初始化 ClusteringWithTopic，接受一个 n_topics_list，其中包含多个聚类数目，
         选取 silhouette_score 最高的结果。
         """
-        embedding_model = SentenceTransformer("nomic-ai/nomic-embed-text-v1", trust_remote_code=True)
+        try:
+            embedding_model = SentenceTransformer("nomic-ai/nomic-embed-text-v1", trust_remote_code=True, cache_folder=cache_dir)
+        except Exception as e:
+            print(f"Error initializing SentenceTransformer: {e}")
+            embedding_model = SentenceTransformer("nomic-ai/nomic-embed-text-v1", trust_remote_code=True)
         self.embeddings = embedding_model.encode(df, show_progress_bar=True)
         
         self.df = df
@@ -92,7 +110,7 @@ class ClusteringWithTopic:
         # 用于存储不同聚类数目的结果
         self.best_n_topics = None
         self.best_labels = None
-        self.best_score = -1 
+        self.best_score = -1
     # def fit_and_get_labels(self, X):
     #     topics, probs = self.topic_model.fit_transform(self.df, self.embeddings)
     #     return topics 

src/demo/survey_generation_pipeline/main.py

@@ -27,6 +27,10 @@ import os
 from markdown_pdf import MarkdownPdf, Section  # Assuming you are using markdown_pdf
 from typing import Any
 import xml.etree.ElementTree as ET
+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()
@@ -286,15 +290,19 @@ class ASG_system:
         self.pipeline = None
 
 
-        # model_id = "meta-llama/Meta-Llama-3.1-8B-Instruct"
-        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",
-        # )
+        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")

src/demo/survey_generator_api.py

@@ -9,6 +9,10 @@ import numpy as np
 from numpy.linalg import norm
 import openai
 from .asg_retriever import Retriever
+from .path_utils import get_path, setup_hf_cache
+
+# 设置 Hugging Face 缓存目录
+cache_dir = setup_hf_cache()
 
 def getQwenClient(): 
     # openai_api_key = os.environ.get("OPENAI_API_KEY")
@@ -506,7 +510,7 @@ Survey Paper Content for "{section_title}":
     response = generateResponse(client, formatted_prompt).strip()
     sentences = re.split(r'(?<=[.!?])\s+', response.strip())
 
-    embedder = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
+    embedder = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2", cache_folder=cache_dir)
     sentence_embeddings = embedder.embed_documents(sentences)
     chunk_texts = [c["content"] for c in citation_data_list]
     chunk_sources = [c["source"] for c in citation_data_list]
@@ -627,7 +631,7 @@ Survey Paper Content for "{section_title}":
                 para_index_map.append(p_idx)
 
     # -- 3. 对所有句子进行向量化嵌入（保持逻辑：一次性处理全文） ---
-    embedder = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
+    embedder = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2", cache_folder=cache_dir)
     sentence_embeddings = embedder.embed_documents(all_sentences)
 
     # -- 4. 对 citation_data_list 做向量化嵌入 ---
@@ -763,25 +767,17 @@ def query_embedding_for_title(
     n_results: int = 1, 
     embedder: HuggingFaceEmbeddings = None
 ):
-    final_context = ""
+    if embedder is None:
+        try:
+            embedder = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2", cache_folder=cache_dir)
+        except Exception as e:
+            print(f"Error initializing embedder: {e}")
+            embedder = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
+    
     retriever = Retriever()
-    title_embedding = embedder.embed_query(title)
-
-    query_result = retriever.query_chroma(
-        collection_name=collection_name, 
-        query_embeddings=[title_embedding], 
-        n_results=n_results
-    )
-    # old
-    # query_result_chunks = query_result["documents"][0]
-    # for chunk in query_result_chunks:
-    #     final_context += chunk.strip() + "//\n"
-
-    # 2025
-    if "documents" in query_result and len(query_result["documents"]) > 0:
-        for chunk in query_result["documents"][0]:
-            final_context += chunk.strip() + "//\n"
-    return final_context
+    query_embeddings = embedder.embed_query(title)
+    query_result = retriever.query_chroma(collection_name=collection_name, query_embeddings=[query_embeddings], n_results=n_results)
+    return query_result
 
 # old
 def generate_context_list(outline, collection_list):
@@ -812,32 +808,32 @@ def generate_context_list(outline, collection_list):
 
 # 2025
 def generate_context_list(outline, collection_list):
-
-    subsections = parse_outline_with_subsections(outline)
-    print("[DEBUG] subsections:", subsections)
-
-    embedder = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
-    context_list_final = []
+    try:
+        embedder = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2", cache_folder=cache_dir)
+    except Exception as e:
+        print(f"Error initializing embedder: {e}")
+        embedder = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
     
-    for level, title in subsections:
-        if title.startswith("3"):
-            cluster_idx = 0
-        elif title.startswith("4"):
-            cluster_idx = 1
-        elif title.startswith("5"):
-            cluster_idx = 2
+    retriever = Retriever()
+    context_list = []
+    
+    for section_title in outline:
+        query_embeddings = embedder.embed_query(section_title)
+        final_context = ""
+        seen_chunks = set()
         
-        context_temp = ""
-        for coll_name in collection_list[cluster_idx]:
-            retrieved_context = query_embedding_for_title(
-                collection_name=coll_name,
-                title=title,
-                n_results=3,
-                embedder=embedder
-            )
-            context_temp += retrieved_context + "\n"
-        context_list_final.append(context_temp)    
-    return context_list_final
+        for collection_name in collection_list:
+            query_result = retriever.query_chroma(collection_name=collection_name, query_embeddings=[query_embeddings], n_results=2)
+            query_result_chunks = query_result["documents"][0]
+            
+            for chunk in query_result_chunks:
+                if chunk not in seen_chunks:
+                    final_context += chunk.strip() + "//\n"
+                    seen_chunks.add(chunk)
+        
+        context_list.append(final_context)
+    
+    return context_list
 
 # 1.8 输入introduction 输出带引用 (collection name) 的introduction
 def introduction_with_citations(
@@ -847,110 +843,65 @@ def introduction_with_citations(
     dynamic_threshold: bool = True,
     diversity_limit: int = 3
 ) -> str:
-    """
-    给已生成好的 Introduction 添加引用，保留原有段落结构与换行数量。
-    :param intro_text: 已生成的引言文本（多段）。
-    :param citation_data_list: 需引用的文献块列表，每项为 {"content": "...", "source": "..."}。
-    :param base_threshold: 基础相似度阈值。
-    :param dynamic_threshold: 是否使用动态阈值 (mean + k*std)。
-    :param diversity_limit: 同一 source 最多引用次数。
-    :return: 带有 [paperName] 引用的 Introduction 文本。
-    """
-
-    # 1. 按原有段落拆分
-    paragraphs = intro_text.split('\n\n')
-    if not paragraphs:
-        return intro_text
-
-    # 2. 逐段落拆分句子，记录每句所属段落编号
-    all_sentences = []
-    para_index_map = []
-    for p_idx, para in enumerate(paragraphs):
-        if not para.strip():
-            # 空段落，直接跳过切句，保持段落分隔
-            continue
-        # 用正则在段落内部按 .!? 分句
-        sentences_in_para = re.split(r'(?<=[.!?])\s+', para)
-        for sent in sentences_in_para:
-            if sent:
-                all_sentences.append(sent)
-                para_index_map.append(p_idx)
-
-    # 如果拆不出任何句子，直接返回
-    if not all_sentences:
-        return intro_text
-
-    # 3. 对所有句子进行 Embedding
-    embedder = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
-    sentence_embeddings = embedder.embed_documents(all_sentences)
-
-    # 4. 对 citation_data_list 里每段文献块进行向量化
+    # 将介绍文本按句子分割
+    sentences = re.split(r'(?<=[.!?])\s+', intro_text.strip())
+    
+    # 初始化 embedder
+    try:
+        embedder = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2", cache_folder=cache_dir)
+    except Exception as e:
+        print(f"Error initializing embedder: {e}")
+        embedder = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
+    
+    # 对句子和引用数据进行向量化
+    sentence_embeddings = embedder.embed_documents(sentences)
     chunk_texts = [c["content"] for c in citation_data_list]
     chunk_sources = [c["source"] for c in citation_data_list]
     chunk_embeddings = embedder.embed_documents(chunk_texts)
-
+    
+    # 计算余弦相似度
     def cosine_sim(a, b):
         return np.dot(a, b) / (norm(a) * norm(b) + 1e-9)
-
-    # 5. 建立句子-引用块相似度矩阵
+    
+    # 构建相似度矩阵
     sim_matrix = []
     for s_emb in sentence_embeddings:
         row = [cosine_sim(s_emb, c_emb) for c_emb in chunk_embeddings]
         sim_matrix.append(row)
     sim_matrix = np.array(sim_matrix)
-
-    # 6. 动态阈值(或固定阈值)
+    
+    # 计算动态阈值
     all_sims = sim_matrix.flatten()
     mean_sim = np.mean(all_sims)
-    std_sim  = np.std(all_sims)
+    std_sim = np.std(all_sims)
     k = 0.5
     threshold = max(base_threshold, mean_sim + k * std_sim) if dynamic_threshold else base_threshold
-
-    # 7. 找出相似度 >= threshold 的 (句子ID, 文献块ID, 相似度) 
+    
+    # 找出候选引用
     candidates = []
-    for i in range(len(all_sentences)):
-        for j in range(len(chunk_embeddings)):
-            if sim_matrix[i, j] >= threshold:
-                candidates.append((i, j, sim_matrix[i, j]))
-
-    # 8. 按相似度降序排列
+    for i, sent in enumerate(sentences):
+        for j, sim in enumerate(sim_matrix[i]):
+            if sim >= threshold:
+                candidates.append((i, j, sim))
+    
+    # 按相似度排序并分配引用
+    source_count = {s: 0 for s in chunk_sources}
     candidates.sort(key=lambda x: x[2], reverse=True)
-
-    # 记录：句子 -> 已分配的 source；并限制每个 source 最多引用次数
-    source_count = {src: 0 for src in chunk_sources}
     assigned = {}
-
-    for (sent_id, chk_id, sim_val) in candidates:
+    
+    for (sent_id, chk_id, sim) in candidates:
         if sent_id not in assigned:
             src = chunk_sources[chk_id]
             if source_count[src] < diversity_limit:
                 assigned[sent_id] = src
                 source_count[src] += 1
-
-    # 9. 将引用插入句尾
+    
+    # 更新句子
     updated_sentences = []
-    for i, sentence in enumerate(all_sentences):
+    for i, sentence in enumerate(sentences):
         if i in assigned:
             updated_sentences.append(sentence + f" [{assigned[i]}]")
         else:
             updated_sentences.append(sentence)
-
-    # 10. 按原先段落顺序拼回文本
-    updated_paras = [""] * len(paragraphs)
-    para_sentences_map = [[] for _ in range(len(paragraphs))]
-
-    for s_idx, sent in enumerate(updated_sentences):
-        p_idx = para_index_map[s_idx]
-        para_sentences_map[p_idx].append(sent)
-
-    for i in range(len(paragraphs)):
-        if not paragraphs[i].strip():
-            # 保持空段落不动
-            updated_paras[i] = paragraphs[i]
-        else:
-            # 同段落内的句子用空格拼起来
-            updated_paras[i] = " ".join(para_sentences_map[i])
-
-    # 11. 用原先换行分隔符拼回
-    updated_intro = "\n\n".join(updated_paras)
-    return updated_intro
+    
+    return " ".join(updated_sentences)

src/demo/views.py

@@ -44,7 +44,7 @@ from dotenv import load_dotenv
 from pathlib import Path
 from markdown_pdf import MarkdownPdf, Section
 import tempfile
-from .path_utils import get_path
+from .path_utils import get_path, setup_hf_cache
 
 dotenv_path = os.path.join(os.path.dirname(__file__), ".env")
 load_dotenv()
@@ -59,6 +59,9 @@ load_dotenv()
 # print(f"OPENAI_API_KEY: {openai_api_key}")
 # print(f"OPENAI_API_BASE: {openai_api_base}")
 
+# 设置 Hugging Face 缓存目录
+cache_dir = setup_hf_cache()
+
 # 获取路径配置
 paths_config = get_path('pdf')  # 使用 get_path 函数获取路径配置
 DATA_PATH = get_path('pdf')
@@ -144,8 +147,15 @@ Global_cluster_names = []
 Global_citation_data = []
 Global_cluster_num = 4
 
-
-embedder = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
+try:
+    embedder = HuggingFaceEmbeddings(
+        model_name="sentence-transformers/all-MiniLM-L6-v2",
+        cache_folder=cache_dir
+    )
+except Exception as e:
+    print(f"Error initializing embedder: {e}")
+    # 如果初始化失败，尝试使用默认设置
+    embedder = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
 
 from demo.category_and_tsne import clustering