File size: 2,405 Bytes
28c2a3d 0ee9d00 |
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 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 |
from flask import jsonify, Flask, request
from embedding_model import HuggingfaceSentenceTransformerModel
from similarity_metric import CosineSimilarity
from pdf_parser import GrobidSciPDFPaser
from chatbot import OpenAIChatbot
from chat_pdf import ChatPDF
from config import DEFAULT_ENGINE, MAX_TOKEN_MODEL_MAP, DEFAULT_TEMPERATURE, DEFAULT_TOP_P, DEFAULT_PRESENCE_PENALTY, DEFAULT_FREQUENCY_PENALTY, DEFAULT_REPLY_COUNT
app = Flask(__name__)
chatpdf_pool = {}
embedding_model = HuggingfaceSentenceTransformerModel()
simi_metric = CosineSimilarity()
@app.route("/query/", methods=['POST', 'GET'])
def query():
api_key = request.headers.get('Api-Key')
pdf_link = request.json['pdf_link']
user_stamp = request.json['user_stamp']
user_query = request.json['user_query']
print(
"api_key", api_key,
"pdf_link", pdf_link,
"user_stamp", user_stamp,
"user_query", user_query
)
chat_pdf = None
if user_stamp not in chatpdf_pool:
print(f"User {user_stamp} not in pool, creating new chatpdf")
# Initialize the ChatPDF
bot = OpenAIChatbot(
api_key=api_key,
engine=DEFAULT_ENGINE,
proxy=None,
max_tokens=4000,
temperature=DEFAULT_TEMPERATURE,
top_p=DEFAULT_TOP_P,
presence_penalty=DEFAULT_PRESENCE_PENALTY,
frequency_penalty=DEFAULT_FREQUENCY_PENALTY,
reply_count=DEFAULT_REPLY_COUNT
)
pdf = GrobidSciPDFPaser(
pdf_link=pdf_link
)
chat_pdf = ChatPDF(
pdf=pdf,
bot=bot,
embedding_model=embedding_model,
similarity_metric=simi_metric,
user_stamp=user_stamp
)
chatpdf_pool[user_stamp] = chat_pdf
else:
print("user_stamp", user_stamp, "already exists")
chat_pdf = chatpdf_pool[user_stamp]
try:
response = chat_pdf.chat(user_query)
code = 200
json_dict = {
"code": code,
"response": response
}
except Exception as e:
code = 500
json_dict = {
"code": code,
"response": str(e)
}
return jsonify(json_dict)
# @app.route("/", methods=['GET'])
# def index():
# return "Hello World!"
if __name__ == '__main__':
app.run(host='0.0.0.0', port=5000, debug=False)
|