add requirements.txt, fix fewshot caption generation args

README.md

@@ -30,7 +30,11 @@ At end of this step you should have the following base dirs configured:
 
 ## Installation
 We recommend separate conda environments for retrieval(UniIR) vs generation (UniRAG).
-
+```bash
+conda create -n unirag python=3.10
+conda activate unirag
+pip install -r requirements.txt
+```
 
 ## Environment Variables
 API Keys and other env variables must be in `.env.local`

requirements.txt

@@ -0,0 +1,26 @@
+accelerate>=0.26.0
+apex==0.9.10dev
+diffusers==0.32.2
+ftfy==6.3.1
+huggingface_hub==0.29.2
+jsonlines==4.0.0
+numpy==2.2.3
+openai==1.65.4
+opencv-python==4.11.0.86
+packaging==24.2
+Pillow==11.1.0
+pycocoevalcap==1.2
+python-dotenv==1.0.1
+regex==2024.11.6
+sentencepiece==0.2.0
+setuptools==75.8.2
+spacy==3.8.4
+timm==1.0.15
+torch==2.6.0
+torch-fidelity==0.3.0
+torchmetrics==1.6.2
+torchvision==0.21.0
+tqdm==4.67.1
+transformers==4.49.0
+vertexai==1.71.1
+xformers==0.0.29.post3

src/unirag/caption_generation_inference.py

@@ -377,10 +377,11 @@ def main(args):
             args.base_mbeir_path,
             args.candidates_file_path,
             args.index,
+            args.k,
         )
     else:
         images, retrieval_dict = get_rag_fewshot_examples(
-            args.retrieved_results_path, args.base_mbeir_path, args.index
+            args.retrieved_results_path, args.base_mbeir_path, args.index, args.k
         )
 
     p_class = generator_prompt.Prompt(args.prompt_file, args.k)