GenBankRefs/AI_match_paper.py at main · hivdb/GenBankRefs · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
from dotenv import load_dotenv
from pathlib import Path
from openai import OpenAI, OpenAIError
import os
import time
import pandas as pd


load_dotenv()

TEMPLATE = """
    find the publication for this paper, and just give me the result:
    Title: {Title}
    Authors: {Authors}
    Journal: {Journal}
    Year: {Year}
    Accessions: {Accession}
"""


def build_prompts(system_prompt, metadata):

    prompt = TEMPLATE.format(**metadata)

    prompt = [
        {
            "role": "system",
            "content": system_prompt},
        {
            "role": "user",
            "content": prompt,
        }
    ]

    return prompt


def chat_openai(messages, model='gpt-4o', temperature=0):

    client = OpenAI(api_key=os.getenv('OPENAI_API_KEY'))

    try:
        response = client.responses.create(
            model=model,
            input=messages,
            tools=[{"type": "web_search_preview"}],
            temperature=temperature,
        )
    except OpenAIError as e:
        print(e)
        return []

    resp = [
        i
        for i in response.output
        if 'content' in i.model_fields
    ]
    return resp


def ask_ai(prompt):
    time.sleep(1)
    response = chat_openai(prompt)

    if len(response) == 0:
        return 'No response'
    else:
        return response[0].content[0].text


def using_ai_match(virus, genbank_unmatched, file_suffix, overwrite=False):
    virus_name = virus.full_name if virus.full_name else virus.name

    system_prompt = open(
        Path(__file__).resolve().parent / 'AI_match_template.txt').read().format(
            virus_name=virus_name)

    cache_file = virus.output_excel_dir / f'{virus.name}_{file_suffix}.xlsx'
    answers = []
    if cache_file.exists() and not overwrite:
        answers = pd.read_excel(cache_file)
        answer_map = {
            int(i['RefID']): i['AI_answer']
            for _, i in answers.iterrows()
        }
        for idx, row in genbank_unmatched.iterrows():
            genbank_unmatched.at[idx, 'AI_answer'] = answer_map.get(row['RefID'], '')

    for idx, row in genbank_unmatched.iterrows():
        if 'AI_answer' in row:
            continue
        # if 'Direct Submission' in row['Title']:
        #     continue
        # if 'Patent' in row['Journal']:
        #     continue

        prompt = build_prompts(system_prompt, {
            'Title': row['Title'],
            'Authors': row['Authors'],
            'Journal': row['Journal'] if row['Journal'].lower() != 'unpublished' else '',
            'Year': row['Year'] if row['Year'] else '',
            'Accession': row['accession']
        })
        # print(prompt)
        # raise

        answer = ask_ai(prompt)
        genbank_unmatched.loc[idx, 'AI_answer'] = answer

        answers.append({
            'RefID': row['RefID'],
            'Title': row['Title'],
            'Authors': row['Authors'],
            'Journal': row['Journal'] if row['Journal'].lower() != 'unpublished' else '',
            'Year': row['Year'] if row['Year'] else '',
            'Accession': row['accession'],
            'AI_answer': answer,
        })
        print(idx, 'done')

    pd.DataFrame(answers).to_excel(cache_file, index=False)

    return genbank_unmatched