minor improvements

app.py

@@ -24,6 +24,8 @@ sentence_plan = "1. Introduction sentence\n2. Overview of relevant studies\n3. D
 
 class RequestData(BaseModel):
     abstract: str
+    words: str
+    papers: str
 
 class ResponseData(BaseModel):
     summary: str
@@ -31,7 +33,7 @@ class ResponseData(BaseModel):
 
 @app.post("/generateLiteratureSurvey/", response_model=ResponseData)
 async def generate_literature_survey(request_data: RequestData):
-    summary, ids = summarize(request_data.abstract, llms)
+    summary, ids = summarize(request_data.abstract, request_data.words, request_data.papers, llms)
     return {"summary": summary,
             "ids": ids
             }
@@ -48,11 +50,11 @@ async def root():
       return {"status": 0}
     return {"status": 1}
 
-def summarize(query, llms) :
+def summarize(query, n_words, n_papers, llms) :
    keywords = helper.extract_keywords(llms['feature_extractor'], query)
-   papers = helper.search_papers(llms['arxiv_agent'], keywords)
-   ranked_papers = helper.re_rank_papers(llms['ranker'], query, papers)
-   literature_review, ids = helper.generate_related_work(llms['summarizer'], llms['summarizer_tokenizer'], query, ranked_papers, base_prompt, sentence_plan)
+   papers = helper.search_papers(llms['arxiv_agent'], keywords, int(n_papers)*2)
+   ranked_papers = helper.re_rank_papers(llms['ranker'], query, papers, int(n_papers))
+   literature_review, ids = helper.generate_related_work(llms['summarizer'], llms['summarizer_tokenizer'], query, ranked_papers, base_prompt, sentence_plan, int(n_words))
    return literature_review, ids
 
 print("Program running")

helper.py

@@ -53,14 +53,14 @@ def extract_keywords(model, abstract):
     return keyphrases
 
 
-def search_papers(arxiv_agent, keywords):
+def search_papers(arxiv_agent, keywords, n_papers):
     query = " ".join(keywords)
     results = arxiv_agent.get_summaries_as_docs(query)
     #print("arxiv ouptut ")
     #print(results)
     return results
 
-def re_rank_papers(model, query_abstract, papers):
+def re_rank_papers(model, query_abstract, papers, n_papers):
     summaries = {paper.page_content : {"Title":paper.metadata['Title']} for paper in papers}
     print(summaries)
     target_embeddings = model.encode([query_abstract])
@@ -121,7 +121,7 @@ def generate_refs(papers) :
         i+=1
     return refs, ids
 
-def generate_related_work(model, tokenizer, query_abstract, ranked_papers, base_prompt, sentence_plan):
+def generate_related_work(model, tokenizer, query_abstract, ranked_papers, base_prompt, sentence_plan, n_words):
     input_text = f"Abstract: {query_abstract}\n"
     i = 1
     for key in ranked_papers.keys():
@@ -141,7 +141,7 @@ def generate_related_work(model, tokenizer, query_abstract, ranked_papers, base_
     )
 
     generation_args = { 
-    "max_new_tokens": 1600, 
+    "max_new_tokens": n_words, 
     "return_full_text": False, 
     "temperature": 0.0, 
     "do_sample": False,