You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
53 lines
1.5 KiB
53 lines
1.5 KiB
from fastapi import FastAPI, BackgroundTasks |
|
from pydantic import BaseModel |
|
from typing import Optional |
|
|
|
from prompts import get_summary_prompt |
|
from _llm import LLM |
|
from _arango import ArangoDB |
|
|
|
app = FastAPI() |
|
|
|
|
|
class DocumentData(BaseModel): |
|
text: str |
|
arango_db_name: str |
|
arango_id: str |
|
is_sci: Optional[bool] = False |
|
|
|
|
|
@app.post("/summarise_document") |
|
async def summarize_document(doc_data: DocumentData, background_tasks: BackgroundTasks): |
|
background_tasks.add_task(summarise_document_task, doc_data.dict()) |
|
return {"message": "Document summarization has started."} |
|
|
|
|
|
def summarise_document_task(doc_data: dict): |
|
text = doc_data.get("text") |
|
is_sci = doc_data.get("is_sci", False) |
|
|
|
system_message = "You are summarising scientific articles. It is very important that you keep to what is written and do not add any of your own opinions or interpretations. Always answer in English." |
|
llm = LLM(system_message=system_message) |
|
|
|
summary = llm.generate(query=get_summary_prompt(text, is_sci)) |
|
|
|
summary_doc = { |
|
"text_sum": summary, |
|
"meta": { |
|
"model": llm.model, |
|
"system_message": system_message, |
|
"temperature": llm.options["temperature"], |
|
}, |
|
} |
|
|
|
arango = ArangoDB(db_name=doc_data.get("arango_db_name")) |
|
arango.db.update_document( |
|
{"summary": summary_doc, "_id": doc_data.get("arango_id")}, |
|
silent=True, |
|
check_rev=False, |
|
) |
|
|
|
|
|
if __name__ == "__main__": |
|
import uvicorn |
|
uvicorn.run(app, host="0.0.0.0", port=8100)
|
|
|