Biomanus/biomni_predict.py

import argparse
import json
import re
import sys
from pathlib import Path

sys.path.append("/225040511/project/Hypo_Bio_OS")

from biomni.agent.a1 import A1


def extract_answer(text: str) -> str:
    text = str(text)
    match = re.search(r"<solution>(.*?)</solution>", text, flags=re.S | re.I)
    if match:
        text = match.group(1)
    text = re.sub(r"</?[^>]+>", "", text)
    text = text.strip()
    marker = re.search(r"##\s*(.*)", text)
    if marker:
        text = marker.group(1).strip()
    return text


def load_payload(input_file: Path) -> dict:
    return json.loads(input_file.read_text(encoding="utf-8"))


def write_payload(output_file: Path, payload: dict) -> None:
    output_file.write_text(json.dumps(payload, ensure_ascii=True), encoding="utf-8")


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--input-file", required=True)
    parser.add_argument("--output-file", required=True)
    args = parser.parse_args()

    request = load_payload(Path(args.input_file))
    prompt = request["prompt"]
    agent_kwargs = request.get("agent_kwargs", {})

    agent = A1(**agent_kwargs)
    _, final_message = agent.go(prompt)
    answer = extract_answer(final_message)

    write_payload(
        Path(args.output_file),
        {
            "answer": answer,
            "final_message": str(final_message),
        },
    )


if __name__ == "__main__":
    main()