adjusted number of retrievals and llm inputs

app.py

@@ -4,6 +4,9 @@ import gradio as gr
 from ragatouille import RAGPretrainedModel
 from huggingface_hub import InferenceClient
 
+retrieve_results = 10 
+llm_results = 5 
+
 client = InferenceClient("mistralai/Mixtral-8x7B-Instruct-v0.1")
 
 generate_kwargs = dict(
@@ -36,7 +39,7 @@ def get_prompt_text(question, context, formatted = True):
       return f"<s>" + f"[INST] {sys_instruction} " +  f" {message} [/INST] </s> "
     return f"Context:\n {context} \n Given the following info, take a deep breath and lets think step by step to answer the question: {question}. Cite the titles of your sources when answering.\n\n"
 
-def get_references(question, retriever, k = 10):
+def get_references(question, retriever, k = retrieve_results):
     rag_out = retriever.search(query=question, k=k)
     return rag_out
 
@@ -59,7 +62,7 @@ with gr.Blocks(theme = gr.themes.Soft()) as demo:
           paper_title = f'''### [{title}](https://arxiv.org/abs/{rag_answer['document_id']})\n'''
           paper_abs = rag_answer['content']
           md_text_updated += paper_title + paper_abs +  '\n---------------\n'+ '\n'
-        prompt = get_prompt_text(message, '\n\n'.join(rag_cleaner(out) for out in rag_out))
+        prompt = get_prompt_text(message, '\n\n'.join(rag_cleaner(out) for out in rag_out[:llm_results]))
         return md_text_updated, prompt
 
     def ask_llm(prompt):