| import json |
| import os |
| import shutil |
| import textwrap |
| from pathlib import Path |
|
|
| |
| from smolagents.utils import AgentError |
|
|
|
|
| def serialize_agent_error(obj): |
| if isinstance(obj, AgentError): |
| return {"error_type": obj.__class__.__name__, "message": obj.message} |
| else: |
| return str(obj) |
|
|
|
|
| def get_image_description(file_name: str, question: str, visual_inspection_tool) -> str: |
| prompt = f"""Write a caption of 5 sentences for this image. Pay special attention to any details that might be useful for someone answering the following question: |
| {question}. But do not try to answer the question directly! |
| Do not add any information that is not present in the image.""" |
| return visual_inspection_tool(image_path=file_name, question=prompt) |
|
|
|
|
| def get_document_description(file_path: str, question: str, document_inspection_tool) -> str: |
| prompt = f"""Write a caption of 5 sentences for this document. Pay special attention to any details that might be useful for someone answering the following question: |
| {question}. But do not try to answer the question directly! |
| Do not add any information that is not present in the document.""" |
| return document_inspection_tool.forward_initial_exam_mode(file_path=file_path, question=prompt) |
|
|
|
|
| def get_single_file_description(file_path: str, question: str, visual_inspection_tool, document_inspection_tool): |
| file_extension = file_path.split(".")[-1] |
| if file_extension in ["png", "jpg", "jpeg"]: |
| file_description = f" - Attached image: {file_path}" |
| file_description += ( |
| f"\n -> Image description: {get_image_description(file_path, question, visual_inspection_tool)}" |
| ) |
| return file_description |
| elif file_extension in ["pdf", "xls", "xlsx", "docx", "doc", "xml"]: |
| file_description = f" - Attached document: {file_path}" |
| image_path = file_path.split(".")[0] + ".png" |
| if os.path.exists(image_path): |
| description = get_image_description(image_path, question, visual_inspection_tool) |
| else: |
| description = get_document_description(file_path, question, document_inspection_tool) |
| file_description += f"\n -> File description: {description}" |
| return file_description |
| elif file_extension in ["mp3", "m4a", "wav"]: |
| return f" - Attached audio: {file_path}" |
| else: |
| return f" - Attached file: {file_path}" |
|
|
|
|
| def get_zip_description(file_path: str, question: str, visual_inspection_tool, document_inspection_tool): |
| folder_path = file_path.replace(".zip", "") |
| os.makedirs(folder_path, exist_ok=True) |
| shutil.unpack_archive(file_path, folder_path) |
|
|
| prompt_use_files = "" |
| for root, dirs, files in os.walk(folder_path): |
| for file in files: |
| file_path = os.path.join(root, file) |
| prompt_use_files += "\n" + textwrap.indent( |
| get_single_file_description(file_path, question, visual_inspection_tool, document_inspection_tool), |
| prefix=" ", |
| ) |
| return prompt_use_files |
|
|
|
|
| def get_tasks_to_run(data, total: int, base_filename: Path, tasks_ids: list[int]): |
| f = base_filename.parent / f"{base_filename.stem}_answers.jsonl" |
| done = set() |
| if f.exists(): |
| with open(f, encoding="utf-8") as fh: |
| done = {json.loads(line)["task_id"] for line in fh if line.strip()} |
|
|
| tasks = [] |
| for i in range(total): |
| task_id = int(data[i]["task_id"]) |
| if task_id not in done: |
| if tasks_ids is not None: |
| if task_id in tasks_ids: |
| tasks.append(data[i]) |
| else: |
| tasks.append(data[i]) |
| return tasks |
|
|
