Create summaru.py

generators/summaru.py

@@ -0,0 +1,15 @@
+from rag.retriever import retrieve, format_citations
+from rag.prompts import SUMMARY_SYS, SUMMARY_USER
+from app.deps import gen
+
+def make_summary(query: str):
+    hits = retrieve(query, top_k=10)
+    links = format_citations(hits)  # [(id, url)]
+    link_lines = "\n".join([f"[{i}] {u}" for i,u in links])
+    ctx = "\n".join([f"[chunk:{h['chunk_id']}] {h['text']}" for h in hits])
+
+    prompt = [{"role":"system","content":SUMMARY_SYS},
+              {"role":"user","content":SUMMARY_USER.format(links=link_lines, contexts=ctx)}]
+
+    out = gen(prompt, do_sample=False, temperature=0.2)[0]["generated_text"]
+    return out, links