Spaces:
Sleeping
Sleeping
File size: 1,748 Bytes
9aa78e6 ad329be 9aa78e6 9e36d1d 9aa78e6 ad329be 9aa78e6 ad329be 9aa78e6 d22c38f 9aa78e6 |
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 |
from ..utils.functions import getConfig, getYaml
from ..utils.logger import logger
import litellm
import os
class SummaryEngine:
def __init__(self):
logger.info("INITIALIZING SUMMARY ENGINE")
self.config = getConfig(os.path.join(os.getcwd(), "config.ini"))
self.prompts = getYaml(os.path.join(os.getcwd(), "prompts.yaml"))
def summarize(self, texts: list[str]) -> str:
"""
Summarize a text
Args:
texts: list of texts to summarize
Returns:
summary: summary of the texts
"""
try:
logger.info("Summarizing the details extracted from the images")
allSummaries = str()
for i in range(len(texts)):
batchSize = self.config.getint("DETAIL EXTRACTOR", "BATCHSIZE")
startPage = i * batchSize + 1
endPage = startPage + batchSize - 1
allSummaries += f"# SUMMARY FOR PAGES {startPage}-{endPage}:\n" + texts[i]
completion = litellm.completion(
model = self.config.get("SUMMARIZER", "LLM"),
api_key = os.environ["GROQ_API_KEY"],
api_base = self.config.get("GROQ CONFIG", "BASEURL"),
messages = [
{"role": "system", "content": self.prompts["summaryEnginePrompt"]},
{"role": "user", "content": f"AGGEREGATED SUMMARIES: {allSummaries}"}
]
)
response = completion["choices"][0]["message"]["content"]
logger.info("Summary generated successfully")
return response
except Exception as e:
logger.exception(f"Error summarizing the text: {e}")
return None |