add og prompt

app.py

@@ -37,11 +37,9 @@ class GeneralRetrievalQuery(BaseModel):
     visual_element_explanation: str
 
 
-def get_retrieval_prompt(prompt_name: str) -> Tuple[str, GeneralRetrievalQuery]:
-    if prompt_name != "general":
-        raise ValueError("Only 'general' prompt is available in this version")
-
-    prompt = """You are an AI assistant specialized in document retrieval tasks. Given an image of a document page, your task is to generate retrieval queries that someone might use to find this document in a large corpus.
+def get_retrieval_prompt(prompt_name: str) -> Tuple[str, BaseModel]:
+    if prompt_name == "general":
+        prompt = """You are an AI assistant specialized in document retrieval tasks. Given an image of a document page, your task is to generate retrieval queries that someone might use to find this document in a large corpus.
 
 Please generate 3 different types of retrieval queries:
 
@@ -73,15 +71,52 @@ Here is the document image to analyze:
 <image>
 
 Generate the queries based on this image and provide the response in the specified JSON format."""
+        return prompt, GeneralRetrievalQuery
+    elif prompt_name == "multimodal_rag":
+        prompt = """You are an assistant specialized in Multimodal RAG tasks.
+
+The task is the following: given an image from a pdf page, you will have to generate questions that can be asked by a user to retrieve information from a large documentary corpus.
+
+The question should be relevant to the page, and should not be too specific or too general. The question should be about the subject of the page, and the answer needs to be found in the page.
 
-    return prompt, GeneralRetrievalQuery
+Remember that the question is asked by a user to get some information from a large documentary corpus that contains multimodal data. Generate a question that could be asked by a user without knowing the existence and the content of the corpus.
+
+Generate as well the answer to the question, which should be found in the page. And the format of the answer should be a list of words answering the question.
+
+Generate at most THREE pairs of questions and answers per page as JSON with the following format, answer ONLY using JSON, NOTHING ELSE:
+
+{
+"questions": [
+{
+"question": "XXXXXX",
+"answer": ["YYYYYY"]
+},
+{
+"question": "XXXXXX",
+"answer": ["YYYYYY"]
+},
+{
+"question": "XXXXXX",
+"answer": ["YYYYYY"]
+}
+]
+}
+
+where XXXXXX is the question and ['YYYYYY'] is the corresponding list of answers that could be as long as needed.
+
+Note: If there are no questions to ask about the page, return an empty list. Focus on making relevant questions concerning the page.
+
+Here is the page:"""
+        return prompt, BaseModel
+    else:
+        raise ValueError("Invalid prompt name")
 
 
 # defined like this so we can later add more prompting options
 prompt, pydantic_model = get_retrieval_prompt("general")
 
 
-def _prep_data_for_input(image):
+def _prep_data_for_input(image, prompt):
     messages = [
         {
             "role": "user",
@@ -111,10 +146,10 @@ def _prep_data_for_input(image):
 
 
 @spaces.GPU
-def generate_response(image):
-    inputs = _prep_data_for_input(image)
-    inputs = inputs.to("cuda")
-
+def generate_response(image, prompt_name="general"):
+    prompt, _ = get_retrieval_prompt(prompt_name)
+    inputs = _prep_data_for_input(image, prompt)
+    inputs.to("cuda")
     generated_ids = model.generate(**inputs, max_new_tokens=200)
     generated_ids_trimmed = [
         out_ids[len(in_ids) :]
@@ -156,7 +191,14 @@ examples = [
 
 demo = gr.Interface(
     fn=generate_response,
-    inputs=gr.Image(type="pil"),
+    inputs=[
+        gr.Image(type="pil"),
+        gr.Dropdown(
+            choices=["ColPali paper prompt", "retrieval focused prompt"],
+            value="general",
+            label="Prompt Type",
+        ),
+    ],
     outputs=gr.Json(),
     title=title,
     description=description,