Forrest99 commited on
Commit
90f5392
·
verified ·
1 Parent(s): acf1767

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +29 -0
app.py ADDED
@@ -0,0 +1,29 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from flask import Flask, request, jsonify
2
+ from sentence_transformers import SentenceTransformer, util
3
+
4
+ app = Flask(__name__)
5
+
6
+ # 预定义代码片段
7
+ CODE = [
8
+ """def sort_list(x): return sorted(x)""",
9
+ """def count_above_threshold(elements, threshold=0):...""",
10
+ """def find_min_max(elements):..."""
11
+ ]
12
+
13
+ # 初始化模型(启动时自动下载)
14
+ model = SentenceTransformer("flax-sentence-embeddings/st-codesearch-distilroberta-base")
15
+ code_emb = model.encode(CODE, convert_to_tensor=True)
16
+
17
+ @app.route('/search', methods=['POST'])
18
+ def search():
19
+ query = request.json.get('query', '')
20
+ query_emb = model.encode(query, convert_to_tensor=True)
21
+ hits = util.semantic_search(query_emb, code_emb)[0]
22
+ best = hits[0]
23
+ return jsonify({
24
+ 'code': CODE[best['corpus_id']],
25
+ 'score': float(best['score'])
26
+ })
27
+
28
+ if __name__ == '__main__':
29
+ app.run(host='0.0.0.0', port=5000)