Spaces:
Sleeping
Sleeping
File size: 1,238 Bytes
51cdd95 52f5fde 51cdd95 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 |
import os
# Create the cache directory if it doesn't exist and set permissions
hf_cache_dir = '/app/hf_cache'
os.makedirs(hf_cache_dir, exist_ok=True)
# Set the HF_HOME environment variable
os.environ['HF_HOME'] = hf_cache_dir
# Application code starts here
import logging
from flask import Flask, request, jsonify
from app.search_engine import PromptSearchEngine
app = Flask(__name__)
# Disable parallelism for tokenizers
os.environ['TOKENIZERS_PARALLELISM'] = 'false'
# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
search_engine = PromptSearchEngine()
@app.route('/search', methods=['POST'])
def search():
data = request.get_json()
query = data.get('query')
n = data.get('n', 5)
use_pinecone = data.get('use_pinecone', True)
logger.info(f"Received query: {query} with n: {n} and use_pinecone: {use_pinecone}")
results = search_engine.most_similar(query, n, use_pinecone)
formatted_results = [{'score': score, 'prompt': prompt} for score, prompt in results]
logger.info(f"Returning results: {formatted_results}")
return jsonify(formatted_results)
if __name__ == '__main__':
logger.info("Starting Flask server")
app.run(debug=True)
|