Update app.py

app.py

@@ -19,21 +19,26 @@ nltk.download('stopwords')
 # 从环境变量中获取 hf_token
 hf_token = os.getenv('HF_TOKEN')
 
+
 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)
-def query(texts):
-    response = requests.post(api_url, headers=headers, json={"inputs": texts})
-    if response.status_code == 200:
-        result = response.json()
-        if isinstance(result, list):
-            return result
-        elif 'error' in result:
-            raise RuntimeError("Error from Hugging Face API: " + result['error'])
-    else:
-        raise RuntimeError("Failed to get response from Hugging Face API, status code: " + str(response.status_code))
+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)
+# def query(texts):
+#     response = requests.post(api_url, headers=headers, json={"inputs": texts})
+#     if response.status_code == 200:
+#         result = response.json()
+#         if isinstance(result, list):
+#             return result
+#         elif 'error' in result:
+#             raise RuntimeError("Error from Hugging Face API: " + result['error'])
+#     else:
+#         raise RuntimeError("Failed to get response from Hugging Face API, status code: " + str(response.status_code))
 
 # 加载嵌入向量数据集
 faqs_embeddings_dataset = load_dataset('chenglu/hf-blogs-baai-embeddings')
@@ -71,7 +76,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 = 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]):