update

app.py

@@ -8,7 +8,7 @@ es = Elasticsearch(hosts=["https://data.neuralnoise.com:9200"],
                    basic_auth=('elastic', os.environ['ES_PASSWORD']),
                    verify_certs=False, ssl_show_warn=False)
 
-def search(query, index="pubmed", num_results=3):
+def search(query, index="pubmed", num_docs=3):
     """
     Search the Elasticsearch index for the most relevant documents.
     """
@@ -20,7 +20,7 @@ def search(query, index="pubmed", num_results=3):
             "match": {
                 "content": query  # Assuming documents have a 'content' field
             }
-        }, "size": num_results
+        }, "size": num_docs
     }
 
     response = es.options(request_timeout=60).search(index=index, body=es_request_body)
@@ -32,15 +32,17 @@ def search(query, index="pubmed", num_results=3):
 
     return docs
 
-def rag_pipeline(prompt, index="pubmed", model_name="HuggingFaceH4/zephyr-7b-beta"):
+def rag_pipeline(prompt, index="pubmed", num_docs=3, model_name="HuggingFaceH4/zephyr-7b-beta"):
     """
     A simple RAG pipeline that retrieves documents and uses them to enrich the context for the LLM.
     """
     # Load your language model from HuggingFace Transformers
     generator = pipeline("text-generation", model=model_name)
 
+    num_docs = int(num_docs)
+
     # Retrieve documents
-    docs = search(prompt, index=index)
+    docs = search(prompt, index=index, num_docs=num_docs)
     joined_docs = '\n\n'.join(docs)
     
     messages = [
@@ -67,6 +69,7 @@ iface = gr.Interface(fn=rag_pipeline,
                      inputs=[
                          gr.Textbox(label="Input Prompt", value="Are group 2 innate lymphoid cells (ILC2s) increased in chronic rhinosinusitis with nasal polyps or eosinophilia?"),
                          gr.Dropdown(label="Index", choices=["pubmed", "wikipedia", "textbooks"], value="pubmed"),
+                         gr.Number(label="Number of Documents", value=3, step=1, min_value=1, max_value=10),
                          gr.Dropdown(label="Model", choices=["HuggingFaceH4/zephyr-7b-beta", "meta-llama/Llama-2-7b-chat-hf", "meta-llama/Llama-2-13b-chat-hf", "meta-llama/Llama-2-70b-chat-hf"], value="HuggingFaceH4/zephyr-7b-beta")
                      ],
                      outputs=[