|
| 1 | +""" |
| 2 | +ADE (Adverse Drug Event) probability estimation task: |
| 3 | +
|
| 4 | +Given a pair of (Drug Category, Adverse Event), have the LLM generate an estimate |
| 5 | +of the probability that the drug category is associated with an increased risk |
| 6 | +of the adverse event. |
| 7 | +
|
| 8 | +Run this N times (without caching) to get statistics on the estimates. |
| 9 | +Illustrates the use of `llm_response_batch`. |
| 10 | +
|
| 11 | +Default model is GPT4o, see how to specify alternative models below. |
| 12 | +
|
| 13 | +Example run: |
| 14 | +
|
| 15 | +python3 examples/basic/ drug-outcomes.py \ |
| 16 | + --model litellm/claude-3-5-sonnet-20240620 --temp 0.1 \ |
| 17 | + --pair "(Antibiotics, Acute Liver Injury)" --n 20 --reason true |
| 18 | +
|
| 19 | +Interesting models to try: |
| 20 | +- gpt-4o (default) |
| 21 | +- gpt-4 |
| 22 | +- litellm/claude-3-5-sonnet-20240620 |
| 23 | +- groq/llama3-70b-8192 |
| 24 | +
|
| 25 | +See reference below for specific (DrugCategory, ADE) pairs to test. |
| 26 | +
|
| 27 | +References: |
| 28 | + - Guides to using Langroid with local and non-OpenAI models: |
| 29 | + https://langroid.github.io/langroid/tutorials/local-llm-setup/ |
| 30 | + https://langroid.github.io/langroid/tutorials/non-openai-llms/ |
| 31 | + - OMOP Ground Truth table of known Drug-ADE associations: |
| 32 | + (see page 16 for the table of Drug-ADE pairs) |
| 33 | + https://www.brookings.edu/wp-content/uploads/2012/04/OMOP-methods-review.pdf |
| 34 | +""" |
| 35 | + |
| 36 | +import langroid as lr |
| 37 | +import langroid.language_models as lm |
| 38 | +from langroid.utils.configuration import settings |
| 39 | +import numpy as np |
| 40 | +import re |
| 41 | +from fire import Fire |
| 42 | + |
| 43 | +# Turn off cache retrieval, to get independent estimates on each run |
| 44 | +settings.cache = False |
| 45 | + |
| 46 | +MODEL = lm.OpenAIChatModel.GPT4o |
| 47 | +TEMP = 0.1 |
| 48 | +PAIR = "(Antibiotics, Acute Liver Injury)" |
| 49 | +N = 20 |
| 50 | +# should LLM include reasoning along with probability? |
| 51 | +# (meant to test whether including reasoning along with the probability |
| 52 | +# improves accuracy and/or variance of estimates) |
| 53 | +REASON: bool = False |
| 54 | + |
| 55 | + |
| 56 | +def extract_num(x: str) -> int: |
| 57 | + """ |
| 58 | + Extracts an integer from a string that contains a number. |
| 59 | +
|
| 60 | + Args: |
| 61 | + x (str): The input string containing the number. |
| 62 | +
|
| 63 | + Returns: |
| 64 | + int: The extracted integer. |
| 65 | +
|
| 66 | + Raises: |
| 67 | + ValueError: If no number is found in the expected format. |
| 68 | + """ |
| 69 | + match = re.search(r"\d+", x) |
| 70 | + if match: |
| 71 | + return int(match.group(0)) |
| 72 | + else: |
| 73 | + return -1 |
| 74 | + |
| 75 | + |
| 76 | +def main( |
| 77 | + model: str = MODEL, |
| 78 | + temp: float = TEMP, |
| 79 | + pair: str = PAIR, |
| 80 | + n: int = N, |
| 81 | + reason: bool = REASON, |
| 82 | +): |
| 83 | + REASONING_PROMPT = ( |
| 84 | + """ |
| 85 | + IMPORTANT: Before showing your estimated probability, |
| 86 | + you MUST show 2-3 sentences with your REASONING, and THEN give your |
| 87 | + percent probability estimate in the range [0,100]. |
| 88 | + """ |
| 89 | + if reason |
| 90 | + else "" |
| 91 | + ) |
| 92 | + |
| 93 | + agent = lr.ChatAgent( |
| 94 | + lr.ChatAgentConfig( |
| 95 | + llm=lm.OpenAIGPTConfig( |
| 96 | + temperature=temp, |
| 97 | + chat_model=model, |
| 98 | + ), |
| 99 | + name="ADE-Estimator", |
| 100 | + system_message=f""" |
| 101 | + You are a clinician with deep knowledge of Adverse Drug Events (ADEs) |
| 102 | + of various drugs and categories of drugs. |
| 103 | + You will be given a (DRUG CATEGORY, ADVERSE OUTCOME) pair, |
| 104 | + you have to estimate the probability that this DRUG CATEGORY |
| 105 | + is associated with INCREASED RISK of the ADVERSE OUTCOME. |
| 106 | + |
| 107 | + {REASONING_PROMPT} |
| 108 | + |
| 109 | + You must give your probability estimate as a SINGLE NUMBER e.g. 56, |
| 110 | + which means 56%. |
| 111 | + DO NOT GIVE A RANGE OF PROBABILITIES, ONLY A SINGLE NUMBER. |
| 112 | + """, |
| 113 | + ) |
| 114 | + ) |
| 115 | + |
| 116 | + results = lr.llm_response_batch( |
| 117 | + agent, |
| 118 | + [pair] * n, |
| 119 | + # ["(Beta Blockers, Mortality after Myocardial Infarction)"]*20, |
| 120 | + ) |
| 121 | + probs = [extract_num(r.content) for r in results] |
| 122 | + cached = [r.metadata.cached for r in results] |
| 123 | + n_cached = sum(cached) |
| 124 | + # eliminate negatives (due to errs) |
| 125 | + probs = [p for p in probs if p >= 0] |
| 126 | + mean = np.mean(probs) |
| 127 | + std = np.std(probs) |
| 128 | + std_err = std / np.sqrt(len(probs)) |
| 129 | + hi = max(probs) |
| 130 | + lo = min(probs) |
| 131 | + print(f"Stats for {pair} with {model} temp {temp} reason {reason}:") |
| 132 | + print( |
| 133 | + f"N: {len(probs)} ({n_cached} cached ) Mean: {mean:.2f}, Std: {std:.2f}, StdErr:" |
| 134 | + f" {std_err:.2f}, min: {lo:.2f}, max: {hi:.2f}" |
| 135 | + ) |
| 136 | + toks, cost = agent.llm.tot_tokens_cost() |
| 137 | + print(f"Tokens: {toks}, Cost: {cost:.2f}") |
| 138 | + |
| 139 | + |
| 140 | +if __name__ == "__main__": |
| 141 | + Fire(main) |
0 commit comments