scripts/AutoGen-MASSA.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import time\n",
    "from datetime import timedelta\n",
    "import json\n",
    "import scipy\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from PIL import Image\n",
    "import torch\n",
    "import autogen\n",
    "from sklearn.metrics import accuracy_score, roc_curve, auc, precision_recall_curve,average_precision_score\n",
    "from sklearn.linear_model import LinearRegression\n",
    "\n",
    "import sys\n",
    "sys.path.append('..')\n",
    "sys.path.append('../src')\n",
    "import src.inference as src_inference\n",
    "import src.finetune as src_finetune\n",
    "\n",
    "import seaborn as sns\n",
    "sns.set(context='paper', style='white')\n",
    "sns.set_color_codes()\n",
    "import matplotlib.pyplot as plt\n",
    "# %matplotlib inline\n",
    "\n",
    "import warnings\n",
    "warnings.filterwarnings('ignore')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.Inference Function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def Analyze_Protein_Property(file_path, task_type, finetune=False):\n",
    "    task_type = task_type.lower()\n",
    "    if not finetune:\n",
    "        return src_inference.launch(file_path, task_type)\n",
    "    else:\n",
    "        return src_finetune.launch(file_path, task_type)\n",
    "    \n",
    "def Analyze_Protein_Drug_Interaction(file_path, task_type, finetune=False):\n",
    "    task_type = task_type.lower()\n",
    "    if not finetune:\n",
    "        return src_inference.launch(file_path, task_type)\n",
    "    else:\n",
    "        return src_finetune.launch(file_path, task_type)\n",
    "    \n",
    "def Analyze_Protein_Protein_Interaction(file_path, task_type, finetune=False):\n",
    "    task_type = task_type.lower()\n",
    "    if not finetune:\n",
    "        return src_inference.launch(file_path, task_type)\n",
    "    else:\n",
    "        return src_finetune.launch(file_path, task_type)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.Evaluate Function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def reg_metrics(labels, preds):\n",
    "    labels = np.array(labels) if isinstance(labels, list) else labels\n",
    "    preds = np.array(preds) if isinstance(preds, list) else preds\n",
    "    rmse = ((preds - labels) ** 2).mean() ** 0.5\n",
    "    mae = (np.abs(preds - labels)).mean()\n",
    "    corr = scipy.stats.pearsonr(preds, labels)\n",
    "    lr = LinearRegression()\n",
    "    lr.fit(preds.reshape(-1, 1), labels)\n",
    "    y_ = lr.predict(preds.reshape(-1, 1))\n",
    "    sd = (((labels - y_) ** 2).sum() / (len(labels) - 1)) ** 0.5\n",
    "    return rmse, mae, corr, sd\n",
    "            \n",
    "def Evaluate(inference_file_path, task_type):\n",
    "    task_type = task_type.lower()\n",
    "    inference_dict = json.load(open(inference_file_path, 'r'))\n",
    "    metrics_dict = {}\n",
    "    if task_type in ['stability', 'fluorescence', 'pdbbind']:\n",
    "        for key in inference_dict.keys():\n",
    "            if isinstance(inference_dict[key], list):\n",
    "                labels = [label for _ , (pro_id, label, pred) in enumerate(inference_dict[key])]\n",
    "                preds = [pred for _ , (pro_id, label, pred) in enumerate(inference_dict[key])]\n",
    "                rmse, mae, corr, sd = reg_metrics(labels, preds)\n",
    "                metrics_dict[key] = {'rmse': rmse, 'mae': mae, 'corr': corr, 'sd': sd, 'label': labels, 'pred': preds}\n",
    "    elif 'remote' in task_type or 'homology' in task_type:\n",
    "        for key in inference_dict.keys():\n",
    "            if isinstance(inference_dict[key], list):\n",
    "                labels = [label for _ , (pro_id, label, pred) in enumerate(inference_dict[key])]\n",
    "                preds = [pred for _ , (pro_id, label, pred) in enumerate(inference_dict[key])]\n",
    "                accuracy = accuracy_score(labels, preds)\n",
    "                metrics_dict[key] = {'accuracy': accuracy}\n",
    "    elif 'secondary' in task_type or 'structure' in task_type:\n",
    "        for key in inference_dict.keys():\n",
    "            if isinstance(inference_dict[key], list):\n",
    "                labels = [label for _ , (pro_id, label, pred) in enumerate(inference_dict[key])]\n",
    "                preds = [pred for _ , (pro_id, label, pred) in enumerate(inference_dict[key])]\n",
    "                labels, preds = sum(labels, []), sum(preds, [])\n",
    "                accuracy = accuracy_score(labels, preds)\n",
    "                metrics_dict[key] = {'accuracy': accuracy}\n",
    "    elif task_type in ['kinase'] or ('antigen' in task_type or 'binding' in task_type):\n",
    "        for key in inference_dict.keys():\n",
    "            if isinstance(inference_dict[key], list):\n",
    "                labels = [label for _ , (pro_id, label, pred) in enumerate(inference_dict[key])]\n",
    "                preds = [pred for _ , (pro_id, label, pred) in enumerate(inference_dict[key])]\n",
    "                fpr, tpr, thersholds = roc_curve(labels, preds)\n",
    "                roc_auc = auc(fpr, tpr)\n",
    "                precision, recall, _ = precision_recall_curve(labels, preds)\n",
    "                prc = average_precision_score(labels, preds)\n",
    "                metrics_dict[key] = {'fpr': fpr.tolist(), 'tpr': tpr.tolist(), 'roc_auc': roc_auc, \\\n",
    "                                    'precision': precision.tolist(), 'recall': recall.tolist(), 'prc': prc}\n",
    "    elif task_type in ['skempi']:\n",
    "        for key in inference_dict.keys():\n",
    "            if isinstance(inference_dict[key], list):\n",
    "                labels = [label for _ , (pro_id, label, pred) in enumerate(inference_dict[key])]\n",
    "                preds = [pred for _ , (pro_id, label, pred) in enumerate(inference_dict[key])]\n",
    "                rmse, mae, corr, sd = reg_metrics(labels, preds)\n",
    "                metrics_dict[key] = {'rmse': rmse, 'mae': mae, 'corr': corr, 'sd': sd, 'label': labels, 'pred': preds}\n",
    "            elif isinstance(inference_dict[key], dict):\n",
    "                metrics_dict[key] = {}\n",
    "                for sub_key in results[key]:\n",
    "                    labels = [label for _ , (pro_id, label, pred) in enumerate(inference_dict[key][sub_key])]\n",
    "                    preds = [pred for _ , (pro_id, label, pred) in enumerate(inference_dict[key][sub_key])]\n",
    "                    rmse, mae, corr, sd = reg_metrics(labels, preds)\n",
    "                    metrics_dict[key][sub_key] = {'rmse': rmse, 'mae': mae, 'corr': corr, 'sd': sd, 'label': labels, 'pred': preds}\n",
    "    \n",
    "    metrics_results_dir = \"metrics_results\"\n",
    "    os.makedirs(metrics_results_dir, exist_ok=True)\n",
    "    with open(os.path.join(metrics_results_dir, f\"{task_type}_metrics.json\"), 'w') as f:\n",
    "        json.dump(metrics_dict, f, indent=4)\n",
    "    f.close()\n",
    "    \n",
    "    return os.path.join(metrics_results_dir, f\"{task_type}_metrics.json\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3.Visualize Function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def combined_image(image_files, save_path):\n",
    "    width, height = Image.open(image_files[0]).size\n",
    "    total_width = width * len(image_files)\n",
    "    max_height = height  \n",
    "    new_im = Image.new('RGB', (total_width, max_height))\n",
    "    x_offset = 0  \n",
    "    for img_path in image_files:\n",
    "        img = Image.open(img_path)\n",
    "        new_im.paste(img, (x_offset, 0))\n",
    "        x_offset += img.width\n",
    "    new_im.save(save_path)\n",
    "    for img_path in image_files:\n",
    "        os.remove(img_path)\n",
    "    \n",
    "def plot_reg(visualization_results_dir, metrics_dict, set_colors, task_type, min_value, max_value, interval):\n",
    "    plt.rcParams['font.size'] = 11\n",
    "    plt.rcParams['font.family'] = 'DeJavu Serif'\n",
    "    plt.rcParams['font.serif'] = ['Times New Roman']\n",
    "    plt.rcParams['font.sans-serif'] = ['SimHei']\n",
    "    plt.rcParams['axes.unicode_minus'] = False\n",
    "    \n",
    "    image_files = []\n",
    "    for key, metrics in metrics_dict.items():\n",
    "        set_name = key.capitalize() + (' set (core2016)' if task_type == 'pdbbind' and key == 'test' else ' set')\n",
    "        rmse, mae, corr, sd = metrics['rmse'], metrics['mae'], metrics['corr'], metrics['sd']\n",
    "        table = pd.DataFrame()\n",
    "        table['real'] = metrics['label']\n",
    "        table['predicted'] = metrics['pred']\n",
    "        \n",
    "        grid = sns.jointplot('real', 'predicted', data=table, kind='scatter', color=set_colors[set_name],\n",
    "                             space=0, size=4, ratio=4, s=20, edgecolor='black')\n",
    "        grid.ax_joint.set_xticks(np.arange(min_value, max_value, interval))\n",
    "        grid.ax_joint.set_yticks(np.arange(min_value, max_value, interval))\n",
    "        grid.set_axis_labels(xlabel='real', ylabel='predicted', fontsize=12)\n",
    "        grid.fig.suptitle(task_type.title(), fontsize=14)\n",
    "        grid.fig.subplots_adjust(top=0.92)\n",
    "\n",
    "        if task_type in ['stability', 'fluorescence']:\n",
    "            grid.ax_joint.text(-1.8, 3.5, set_name, fontsize=13)\n",
    "            grid.ax_joint.text(4.3, 5.6, 'Pearson: %.2f ' % corr[0])\n",
    "            grid.ax_joint.text(4.3, 5.1, 'RMSE: %.2f' % (rmse))   \n",
    "            grid.ax_joint.text(4.3, 4.6, 'MAE: %.2f' % (mae))\n",
    "            grid.ax_joint.text(4.3, 4.1, 'SD: %.2f ' % sd)\n",
    "        elif task_type in ['pdbbind']:\n",
    "            grid.ax_joint.text(1, 14, set_name, fontsize=13)\n",
    "            grid.ax_joint.text(16, 19.5, 'Pearson: %.2f ' % corr[0])   \n",
    "            grid.ax_joint.text(16, 18.5, 'RMSE: %.2f' % (rmse))   \n",
    "            grid.ax_joint.text(16, 17.5, 'MAE: %.2f' % (mae))\n",
    "            grid.ax_joint.text(16, 16.5, 'SD: %.2f ' % sd)\n",
    "    \n",
    "        sub_image_file = os.path.join(visualization_results_dir, f'{task_type}_{key}_visualization.jpg')\n",
    "        image_files.append(sub_image_file)\n",
    "        grid.fig.savefig(sub_image_file, bbox_inches='tight', dpi=1024)\n",
    "\n",
    "    combined_image_file = os.path.join(visualization_results_dir, f'{task_type}_visualization.jpg')\n",
    "    combined_image(image_files, combined_image_file)\n",
    "    return combined_image_file\n",
    "\n",
    "\n",
    "def plot_cls(visualization_results_dir, metrics_dict, alternative_colors, task_type):\n",
    "\n",
    "    plt.rcParams['font.size'] = 12\n",
    "    plt.rcParams['font.family'] = 'DeJavu Serif'\n",
    "    plt.rcParams['font.serif'] = ['Times New Roman']\n",
    "    plt.rcParams['font.sans-serif'] = ['SimHei']\n",
    "    plt.rcParams['axes.unicode_minus'] = False\n",
    "\n",
    "    categories, accuracies, colors = [], [], []\n",
    "    for i, (key, metrics) in enumerate(metrics_dict.items()):\n",
    "        categories.append(key.capitalize())\n",
    "        accuracies.append(metrics['accuracy'])\n",
    "        colors.append(alternative_colors[i])\n",
    "\n",
    "    plt.subplots(figsize=(7,4))\n",
    "    bars = plt.bar(categories, accuracies, color=colors)\n",
    "    plt.xticks(categories, rotation=30, fontsize=13)\n",
    "    plt.yticks(np.arange(0, 0.61, 0.1), fontsize=12)\n",
    "    plt.ylim(0, 0.6)\n",
    "    plt.tick_params(axis='both', direction='in', length=6, width=2)\n",
    "    plt.grid(True, axis='y')\n",
    "    for bar in bars:\n",
    "        yval = bar.get_height()\n",
    "        plt.text(bar.get_x() + bar.get_width()/2, yval, round(yval, 3), ha='center', va='bottom', fontsize=13.5)\n",
    "    plt.ylabel('Accuracy', fontsize=13)\n",
    "    plt.title(task_type.title(), fontsize=14)\n",
    "    plt.tight_layout()\n",
    "    image_file = os.path.join(visualization_results_dir, f'{task_type}_visualization.jpg')\n",
    "    plt.savefig(image_file, bbox_inches='tight', dpi=1024)\n",
    "\n",
    "    return image_file\n",
    "\n",
    "def plot_roc_curve(visualization_results_dir, metrics_dict, alternative_colors, task_type):\n",
    "    plt.rcParams['font.size'] = 12\n",
    "    plt.rcParams['font.family'] = 'DeJavu Serif'\n",
    "    plt.rcParams['font.serif'] = ['Times New Roman']\n",
    "    plt.rcParams['font.sans-serif'] = ['SimHei']\n",
    "    plt.rcParams['axes.unicode_minus'] = False\n",
    "    plt.subplots(figsize=(7, 6))\n",
    "     \n",
    "    for i, (key, metrics) in enumerate(metrics_dict.items()):\n",
    "        fpr, tpr, roc_auc = metrics['fpr'], metrics['tpr'], metrics['roc_auc']\n",
    "        plt.plot(fpr, tpr, color=alternative_colors[i], lw=2, label='%s (AUC = %0.2f)' % (key.capitalize(), roc_auc))       \n",
    "    plt.plot([0, 1], [0, 1], color='grey', lw=2, label='Random Guess',linestyle='--')\n",
    "    plt.xlim([0.0, 1.0])\n",
    "    plt.ylim([0.0, 1.05])\n",
    "    plt.xticks(fontsize=12)  \n",
    "    plt.yticks(fontsize=12)  \n",
    "    current_yticks = plt.gca().get_yticks() \n",
    "    for y in current_yticks:\n",
    "        plt.axhline(y=y, color='grey', linestyle='-', lw=0.77, alpha=0.37) \n",
    "    plt.xlabel('False Positive Rate',fontsize=13)\n",
    "    plt.ylabel('True Positive Rate',fontsize=13)\n",
    "    plt.title('ROC curves on Kinase',fontsize=14)\n",
    "    plt.legend(loc=\"lower right\",fontsize=14)\n",
    "    plt.tight_layout()\n",
    "    image_file = os.path.join(visualization_results_dir, f'{task_type}_roc_curve_visualization.jpg')\n",
    "    plt.savefig(image_file, bbox_inches='tight', dpi=1024)\n",
    "    \n",
    "    return image_file\n",
    "\n",
    "def plot_prc_curve(visualization_results_dir, metrics_dict, alternative_colors, task_type):\n",
    "    plt.rcParams['font.size'] = 12\n",
    "    plt.rcParams['font.family'] = 'DeJavu Serif'\n",
    "    plt.rcParams['font.serif'] = ['Times New Roman']\n",
    "    plt.rcParams['font.sans-serif'] = ['SimHei']\n",
    "    plt.rcParams['axes.unicode_minus'] = False\n",
    "    plt.subplots(figsize=(7,6))\n",
    "    \n",
    "    for i, (key, metrics) in enumerate(metrics_dict.items()):\n",
    "        recall, precision, prc = metrics['recall'], metrics['precision'], metrics['prc']\n",
    "        plt.plot(recall, precision, color=alternative_colors[i], lw=2, label='%s (PRC = %0.2f)' % (key.capitalize(), prc))\n",
    "    \n",
    "    plt.plot([0, 1], [1, 0], color='grey', lw=2, label='Random Guess',linestyle='--')\n",
    "    plt.xlim([0.0, 1.0])\n",
    "    plt.ylim([0.0, 1.05])\n",
    "    plt.xticks(fontsize=12)  \n",
    "    plt.yticks(fontsize=12)\n",
    "    current_yticks = plt.gca().get_yticks()\n",
    "    for y in current_yticks:\n",
    "        plt.axhline(y=y, color='grey', linestyle='-', lw=0.77, alpha=0.37)\n",
    "    plt.xlabel('Recall',fontsize=13)\n",
    "    plt.ylabel('Precision',fontsize=13)\n",
    "    plt.title('Precision-Recall curves on Kinase',fontsize=14)\n",
    "    plt.legend(loc='lower right',fontsize=14)\n",
    "    plt.tight_layout()\n",
    "    image_file = os.path.join(visualization_results_dir, f'{task_type}_prc_curve_visualization.jpg')\n",
    "    plt.savefig(image_file, bbox_inches='tight', dpi=1024)\n",
    "    \n",
    "    return image_file\n",
    "\n",
    "\n",
    "def plot_binarycls(visualization_results_dir, metrics_dict, alternative_colors, task_type):\n",
    "    roc_image_file = plot_roc_curve(visualization_results_dir, metrics_dict, alternative_colors, task_type)\n",
    "    pr_image_file = plot_prc_curve(visualization_results_dir, metrics_dict, alternative_colors, task_type)\n",
    "    image_files = [roc_image_file, pr_image_file]\n",
    "    combined_image_file = os.path.join(visualization_results_dir, f'{task_type}_visualization.jpg')\n",
    "    combined_image(image_files, combined_image_file)\n",
    "    \n",
    "    return combined_image_file\n",
    "\n",
    "def plot_box(visualization_results_dir, metrics_dict, task_type):\n",
    "    plt.rcParams['font.size'] = 12\n",
    "    plt.rcParams['font.family'] = 'DeJavu Serif'\n",
    "    plt.rcParams['font.serif'] = ['Times New Roman']\n",
    "    plt.rcParams['font.sans-serif'] = ['SimHei']  # 用来正常显示中文标签\n",
    "    plt.rcParams['axes.unicode_minus'] = False  # 用来正常显示负号\n",
    "    \n",
    "    fig, axes = plt.subplots(1, 3, figsize=(12, 4))\n",
    "    box_width = 0.5\n",
    "\n",
    "    results = {}\n",
    "    results['pearson'], results['rmse'], results['mae'] = [], [], []\n",
    "    for fold_result in metrics_dict['test'].values():\n",
    "        results['pearson'].append(fold_result['corr'][0])\n",
    "        results['rmse'].append(fold_result['rmse'])\n",
    "        results['mae'].append(fold_result['mae'])\n",
    "        \n",
    "    sns.boxplot(data=results['pearson'], ax=axes[0], color='skyblue', width=box_width, showfliers=True)\n",
    "    sns.swarmplot(data=results['pearson'], ax=axes[0], color='black', size=5, alpha=0.77)\n",
    "    axes[0].set_title('10 folds Cross-Validation', fontsize=13)\n",
    "    axes[0].set_ylabel(\"Pearson\", fontsize=12)\n",
    "    axes[0].grid(False)\n",
    "    axes[0].set_xticks([])\n",
    "    axes[0].axhline(y=np.median(results['pearson']), xmin=0, xmax=0.25, color='gray', linestyle='--')\n",
    "    axes[0].text(-0.37, np.median(results['pearson']), f\"{np.median(results['pearson']):.3f}\", fontsize=11, color='black', ha='center', va='bottom')\n",
    "\n",
    "    sns.boxplot(data=results['rmse'], ax=axes[1], color='salmon', width=box_width, showfliers=True)\n",
    "    sns.swarmplot(data=results['rmse'], ax=axes[1], color='black', size=5, alpha=0.77)\n",
    "    axes[1].set_title('10 folds Cross-Validation', fontsize=13)\n",
    "    axes[1].set_ylabel('RMSE', fontsize=12)\n",
    "    axes[1].grid(False)\n",
    "    axes[1].set_xticks([])\n",
    "    axes[1].axhline(y=np.median(results['rmse']), xmin=0, xmax=0.25, color='gray', linestyle='--')\n",
    "    axes[1].text(-0.37, np.median(results['rmse']), f\"{np.median(results['rmse']):.4f}\", fontsize=11, color='black', ha='center', va='bottom')\n",
    "\n",
    "    sns.boxplot(data=results['mae'], ax=axes[2], color='lightgreen', width=box_width, showfliers=True)\n",
    "    sns.swarmplot(data=results['mae'], ax=axes[2], color='black', size=5, alpha=0.77)\n",
    "    axes[2].set_title('10 folds Cross-Validation', fontsize=13)\n",
    "    axes[2].set_ylabel('MAE', fontsize=12)\n",
    "    axes[2].grid(False)\n",
    "    axes[2].set_xticks([])\n",
    "    axes[2].axhline(y=np.median(results['mae']), xmin=0, xmax=0.25, color='gray', linestyle='--')\n",
    "    axes[2].text(-0.37, np.median(results['mae']), f\"{np.median(results['mae']):.3f}\", fontsize=11, color='black', ha='center', va='bottom')\n",
    "    \n",
    "    fig.suptitle(task_type.title(), fontsize=14, y=0.97)\n",
    "    plt.tight_layout()\n",
    "    image_file = os.path.join(visualization_results_dir, f'{task_type}_visualization.jpg')\n",
    "    plt.savefig(image_file, bbox_inches='tight', dpi=1024)\n",
    "\n",
    "    return image_file\n",
    "\n",
    "def Visualize(metrics_file_path, task_type):\n",
    "    task_type = task_type.lower()\n",
    "    visualization_results_dir = \"visualization_results\"\n",
    "    os.makedirs(visualization_results_dir, exist_ok=True)\n",
    "    metrics_dict = json.load(open(metrics_file_path, 'r'))\n",
    "    if task_type in ['stability', 'fluorescence']:\n",
    "        set_colors = {'Train set':'red','Valid set': 'slategrey', 'Test set': 'steelblue'}\n",
    "        visualization_results = plot_reg(visualization_results_dir, metrics_dict, set_colors, task_type, min_value=-2, max_value=5, interval=1)\n",
    "    elif task_type in ['pdbbind']:\n",
    "        set_colors = {'Train set': 'blueviolet', 'Valid set': '#5d3954', 'Test set (core2016)': 'darkcyan'}\n",
    "        visualization_results = plot_reg(visualization_results_dir, metrics_dict, set_colors, task_type, min_value=-0, max_value=16, interval=5)\n",
    "    elif ('remote' in task_type or 'homology' in task_type) or ('secondary' in task_type or 'structure' in task_type):\n",
    "        alternative_colors = ['#d4af37', '#8a9a5b', '#536878']\n",
    "        visualization_results = plot_cls(visualization_results_dir, metrics_dict, alternative_colors, task_type)\n",
    "    elif task_type in ['kinase'] or ('antigen' in task_type or 'binding' in task_type):\n",
    "        alternative_colors = ['navy', 'r', 'b']\n",
    "        visualization_results = plot_binarycls(visualization_results_dir, metrics_dict, alternative_colors, task_type)\n",
    "    elif task_type in ['skempi']:\n",
    "        visualization_results = plot_box(visualization_results_dir, metrics_dict, task_type)\n",
    "        \n",
    "    return visualization_results"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### OpenAI key"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "config_list = [{ 'model': 'gpt-4-32k', 'api_key': 'openai-key'}] \n",
    "\n",
    "manager_llm_config = {\n",
    "    \"config_list\": config_list,\n",
    "    \"seed\": 42,\n",
    "}\n",
    "\n",
    "llm_config = {\n",
    "    \"seed\": 42,\n",
    "    \"config_list\": config_list,\n",
    "    \"functions\": [\n",
    "        {\n",
    "            \"name\": \"Analyze_Protein_Property\",\n",
    "            \"description\": \"Analyze the protein properties including stability, fluoresence, remote homology, secondary structure and so on. Finally return the file path of task results\",\n",
    "            \"parameters\": {\n",
    "                \"type\": \"object\",\n",
    "                \"properties\": {\n",
    "                    \"file_path\": {\n",
    "                        \"type\": \"string\",\n",
    "                        \"description\": \"The file path of protein data used to analyze protein properties\"\n",
    "                    },\n",
    "                    \"task_type\": {\n",
    "                        \"type\": \"string\",\n",
    "                        \"description\": \"The task name\"\n",
    "                    },\n",
    "                    \"finetune\": {\n",
    "                        \"type\": \"boolean\",\n",
    "                        \"description\": \"Is fine-tuning needed for the protein task? If the task instructions do not mention finetuning, the default setting is boolean false.\",\n",
    "                        \"default\": False\n",
    "                    }\n",
    "                }\n",
    "            },\n",
    "            \"required\": [\"file_path\", \"task_type\", \"finetune\"]\n",
    "        },\n",
    "        {\n",
    "            \"name\": \"Analyze_Protein_Drug_Interaction\", \n",
    "            \"description\": \"Analyze the interaction of protein and drug including pdbbind, kinase and so on. Finally return the file path of task results\",\n",
    "            \"parameters\": {\n",
    "                \"type\": \"object\",\n",
    "                \"properties\": {\n",
    "                    \"file_path\": {\n",
    "                        \"type\": \"string\",\n",
    "                        \"description\": \"The file path of protein data used to analyze the interaction of protein and drug\"\n",
    "                    },\n",
    "                    \"task_type\": {\n",
    "                        \"type\": \"string\",\n",
    "                        \"description\": \"The task name\"\n",
    "                    },\n",
    "                    \"finetune\": {\n",
    "                        \"type\": \"boolean\",\n",
    "                        \"description\": \"Is fine-tuning needed for the protein task? If the task instructions do not mention finetuning, the default setting is boolean false.\",\n",
    "                        \"default\": False\n",
    "                    }\n",
    "                }\n",
    "            },\n",
    "            \"required\": [\"file_path\", \"task_type\", \"finetune\"]\n",
    "        },\n",
    "        {\n",
    "            \"name\": \"Analyze_Protein_Protein_Interaction\", \n",
    "            \"description\": \"Analyze the interaction of protein and protein including skempi, etc. Finally return the file path of task results\",\n",
    "            \"parameters\": {\n",
    "                \"type\": \"object\",\n",
    "                \"properties\": {\n",
    "                    \"file_path\": {\n",
    "                        \"type\": \"string\",\n",
    "                        \"description\": \"The file path of protein data used to analyze the interaction of protein and protein\"\n",
    "                    },\n",
    "                    \"task_type\": {\n",
    "                        \"type\": \"string\",\n",
    "                        \"description\": \"The task name\"\n",
    "                    },\n",
    "                    \"finetune\": {\n",
    "                        \"type\": \"boolean\",\n",
    "                        \"description\": \"Is fine-tuning needed for the protein task? If the task instructions do not mention finetuning, the default setting is boolean false.\",\n",
    "                        \"default\": False\n",
    "                    }\n",
    "                }\n",
    "            },\n",
    "            \"required\": [\"file_path\", \"task_type\", \"finetune\"]\n",
    "        },\n",
    "        {\n",
    "            \"name\": \"Evaluate\",\n",
    "            \"description\": \"Calculate the metrics of the returned inference results by the tasks including analyzing protein property, analyzing protein-Drug interaction and analyzing protein-protein interaction. Finally return the metrics\",\n",
    "            \"parameters\": {\n",
    "                \"type\": \"object\",\n",
    "                \"properties\": {\n",
    "                \"inference_file_path\": {\n",
    "                        \"type\": \"string\",\n",
    "                        \"description\": \"The path of returned inference results by the tasks\"\n",
    "                    },\n",
    "                    \"task_type\": {\n",
    "                        \"type\": \"string\",\n",
    "                        \"description\": \"The task name\"\n",
    "                    }\n",
    "                }\n",
    "            },\n",
    "            \"required\": [\"inference_file_path\", \"task_type\"]\n",
    "        },\n",
    "        {\n",
    "            \"name\": \"Visualize\",\n",
    "            \"description\": \"Visualize by the metrics and save it as a jpg\",\n",
    "            \"parameters\": {\n",
    "                \"type\": \"object\",\n",
    "                \"properties\": {\n",
    "                    \"metrics_file_path\": {\n",
    "                        \"type\": \"string\",\n",
    "                        \"description\": \"The path of returned metrics results by the tasks\"\n",
    "                    },\n",
    "                   \"task_type\": {\n",
    "                        \"type\": \"string\",\n",
    "                        \"description\": \"The task name\"\n",
    "                    }\n",
    "                }\n",
    "            },\n",
    "            \"required\": [\"metrics_file_path\", \"task_type\"]\n",
    "        }\n",
    "    ]\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Agents"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create a UserProxyAgent instance named \"user_proxy\"\n",
    "user_proxy = autogen.UserProxyAgent(\n",
    "    name=\"user_proxy\",\n",
    "    human_input_mode=\"NEVER\",\n",
    "    max_consecutive_auto_reply=2,\n",
    "    code_execution_config={\"work_dir\": \".\", \"use_docker\": False},\n",
    "    system_message=\"\"\"Reply TERMINATE if the task has been solved at full satisfaction.\n",
    "                    Otherwise, reply CONTINUE, or the reason why the task is not solved yet.\"\"\"\n",
    ")\n",
    "\n",
    "inference_agent = autogen.AssistantAgent(\n",
    "    name=\"Inference\",\n",
    "    llm_config=llm_config,\n",
    "    function_map={\"Analyze_Protein_Property\": Analyze_Protein_Property,\n",
    "                  \"Analyze_Protein_Drug_Interaction\": Analyze_Protein_Drug_Interaction,\n",
    "                  \"Analyze_Protein_Protein_Interaction\": Analyze_Protein_Protein_Interaction\n",
    "                  }\n",
    ")\n",
    "\n",
    "metrics_agent = autogen.AssistantAgent(\n",
    "    name=\"Evaluation\",\n",
    "    llm_config=llm_config,\n",
    "    function_map={\"Evaluate\": Evaluate}\n",
    ")\n",
    "\n",
    "draw_agent = autogen.AssistantAgent(\n",
    "    name=\"Visualization\",\n",
    "    llm_config=llm_config,\n",
    "    function_map={\"Visualize\": Visualize}\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "groupchat = autogen.GroupChat(agents=[user_proxy, inference_agent, metrics_agent, draw_agent], messages=[],max_round=150) # 群聊 最大的chat轮次12次\n",
    "manager = autogen.GroupChatManager(groupchat=groupchat, llm_config=manager_llm_config)  #管理员，可以与之交互"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### WO Finetune"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "* property"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "user_proxy.initiate_chat(manager, message=\"\"\"Can you analyze on stability task which belongs to protein property prediction benchmark, where the data file and format are as follows, and evaluate the predictions based on the task, finally visualize the evaluation results?\n",
    "\n",
    "../downstream_task_small/stability/sequence_go.txt\n",
    "\n",
    "Start the work now.\n",
    "\"\"\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "* PDI"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "user_proxy.initiate_chat(manager, message=\"\"\"Can you analyze on kinase task which belongs to protein-drug interaction prediction benchmark, where the data file and format are as follows, and evaluate the predictions based on the task, finally visualize the evaluation results?\n",
    "\n",
    "downstream_task/kinase/samples_seq_mole_go.txt\n",
    "\n",
    "Start the work now.\n",
    "\"\"\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### With Finetune"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "* property"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "user_proxy.initiate_chat(manager, message=\"\"\"Can you fine-tune and analyze on antigen binding task which belongs to protein property prediction benchmark, where the data file and format are as follows, and evaluate the predictions based on the task, finally visualize the evaluation results?\n",
    "\n",
    "downstream_task/antigen_binding/sequence_go.txt\n",
    "\n",
    "Start the work now.\n",
    "\"\"\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "lavin",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}