Update app.py

app.py

@@ -26,12 +26,22 @@ hf_token = os.getenv('HF_TOKEN')
 model_id = "BAAI/bge-large-en-v1.5"
 feature_extraction_pipeline = pipeline("feature-extraction", model=model_id)
 
-
 # model_id = "BAAI/bge-large-en-v1.5"
 # api_url = f"https://api-inference.huggingface.co/pipeline/feature-extraction/{model_id}"
 # headers = {"Authorization": f"Bearer {hf_token}"}
 
-# @retry(tries=3, delay=10)
+@retry(tries=3, delay=10)
+def query(texts):
+    # 使用特征提取管道获取特征
+    features = feature_extraction_pipeline(texts)
+    
+    # 将特征降维成二维张量（如果它们不是）
+    # 假设 features 是一个列表，每个元素是一个句子的特征
+    embeddings = [torch.tensor(f).mean(dim=0) for f in features]
+    embeddings = torch.stack(embeddings)
+
+    return embeddings
+
 # def query(texts):
 #     response = requests.post(api_url, headers=headers, json={"inputs": texts})
 #     if response.status_code == 200:
@@ -79,8 +89,8 @@ def get_tags_for_local(dataset, local_value):
 def gradio_query_interface(input_text):
     cleaned_text = clean_content(input_text)
     no_stopwords_text = remove_stopwords(cleaned_text)
-    # new_embedding = query(no_stopwords_text)
-    new_embedding = feature_extraction_pipeline(input_text)
+    new_embedding = query(no_stopwords_text)
+    # new_embedding = feature_extraction_pipeline(input_text)
     query_embeddings = torch.FloatTensor(new_embedding)
     hits = util.semantic_search(query_embeddings, dataset_embeddings, top_k=5)
     if all(hit['score'] < 0.6 for hit in hits[0]):