Update app.py

app.py

@@ -1,29 +1,35 @@
 from flask import Flask, request, jsonify
 from sentence_transformers import SentenceTransformer, util
+import os
 
 app = Flask(__name__)
 
 # 预定义代码片段
 CODE = [
     """def sort_list(x): return sorted(x)""",
-    """def count_above_threshold(elements, threshold=0):...""",
-    """def find_min_max(elements):..."""
+    # 其他代码片段...
 ]
 
-# 初始化模型（启动时自动下载）
-model = SentenceTransformer("flax-sentence-embeddings/st-codesearch-distilroberta-base")
+# 初始化模型（显示指定缓存路径）
+model = SentenceTransformer(
+    "flax-sentence-embeddings/st-codesearch-distilroberta-base",
+    cache_folder=os.getenv("TRANSFORMERS_CACHE")
+)
 code_emb = model.encode(CODE, convert_to_tensor=True)
 
 @app.route('/search', methods=['POST'])
 def search():
-    query = request.json.get('query', '')
-    query_emb = model.encode(query, convert_to_tensor=True)
-    hits = util.semantic_search(query_emb, code_emb)[0]
-    best = hits[0]
-    return jsonify({
-        'code': CODE[best['corpus_id']],
-        'score': float(best['score'])
-    })
+    try:
+        query = request.json['query']
+        query_emb = model.encode(query, convert_to_tensor=True)
+        hits = util.semantic_search(query_emb, code_emb)[0]
+        best = hits[0]
+        return jsonify({
+            'code': CODE[best['corpus_id']],
+            'score': float(best['score'])
+        })
+    except Exception as e:
+        return jsonify({"error": str(e)}), 500
 
 if __name__ == '__main__':
     app.run(host='0.0.0.0', port=5000)