File size: 2,408 Bytes
b8f6efa |
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 |
from elasticsearch import Elasticsearch
import os
import json
import requests
ES_URL = os.environ["ES_URL"]
ES_USER = os.environ["ES_USER"]
ES_PASS = os.environ["ES_PASS"]
ES_CA_CERT = os.environ["ES_CA_CERT"]
class ESGPT:
def __init__(self, index_name):
self.es = Elasticsearch(ES_URL, http_auth=(ES_USER, ES_PASS),
ca_certs=ES_CA_CERT, verify_certs=True)
self.index_name = index_name
self.model_engine = os.environ["OPENAI_GPT_ENGINE"]
self.api_key = os.environ["OPENAI_API_KEY"]
def index(self, doc_id, doc):
self.es.index(index=self.index_name,
id=doc_id,
document=doc)
def search(self, query):
body = {
"query": {
"query_string": {"query": query}
}
}
results = self.es.search(index=self.index_name, body=body)
return results['hits']['hits']
def _paper_results_to_text(self, results):
text_result = ""
for paper in results:
title = ""
if "title" in paper["_source"]:
title = paper["_source"]["title"]
abstract = ""
if "abctract" in paper["_source"]:
abstract = paper["_source"]["abstract"]
paper_str = f"{title}:\n{abstract[:100]}\n\n"
text_result += paper_str
return text_result
def summarize(self, query, results):
# Generate summaries for each search result
result_json_str = self._paper_results_to_text(results)
if result_json_str == "":
result_json_str = "No results found"
print(result_json_str[:500])
body = {
"model": self.model_engine,
"prompt": f"Please summarize the following search results for query: {query}:\n{result_json_str[:1000]}",
"max_tokens": 1000,
"n": 1,
"stop": None,
"temperature": 0.5,
"stream": True,
}
headers = {"Content-Type": "application/json",
"Authorization": f"Bearer {self.api_key}"}
resp = requests.post("https://api.openai.com/v1/completions",
headers=headers,
data=json.dumps(body),
stream=True)
return resp
|