G3/aggregate_llm_predictions.py at main · Applied-Machine-Learning-Lab/G3 · 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
import pandas as pd
import re
import ast
from tqdm import tqdm

df_raw = pd.read_csv('./data/im2gps3k/im2gps3k_places365.csv')
zs_df = pd.read_csv('./data/im2gps3k/llm_predict_results_zs.csv')
rag_5_df = pd.read_csv('./data/im2gps3k/5_llm_predict_results_rag.csv')
rag_10_df = pd.read_csv('./data/im2gps3k/10_llm_predict_results_rag.csv')
rag_15_df = pd.read_csv('./data/im2gps3k/15_llm_predict_results_rag.csv')

pattern = r'[-+]?\d+\.\d+'

for i in tqdm(range(zs_df.shape[0])):
    response = zs_df.loc[i, 'response']
    response = ast.literal_eval(response)
    for idx, content in enumerate(response):
        try:
            match = re.findall(pattern, content)
            latitude = match[0]
            longitude = match[1]
            df_raw.loc[i, f'zs_{idx}_latitude'] = latitude
            df_raw.loc[i, f'zs_{idx}_longitude'] = longitude
        except:
            df_raw.loc[i, f'zs_{idx}_latitude'] = '0.0'
            df_raw.loc[i, f'zs_{idx}_longitude'] = '0.0'

for i in tqdm(range(df_raw.shape[0])):
    response = rag_5_df.loc[i, 'rag_response']
    response = ast.literal_eval(response)
    for idx, content in enumerate(response):
        try:
            match = re.findall(pattern, content)
            latitude = match[0]
            longitude = match[1]
            df_raw.loc[i, f'5_rag_{idx}_latitude'] = latitude
            df_raw.loc[i, f'5_rag_{idx}_longitude'] = longitude
        except:
            df_raw.loc[i, f'5_rag_{idx}_latitude'] = '0.0'
            df_raw.loc[i, f'5_rag_{idx}_longitude'] = '0.0'

for i in tqdm(range(df_raw.shape[0])):
    response = rag_10_df.loc[i, 'rag_response']
    response = ast.literal_eval(response)
    for idx, content in enumerate(response):
        try:
            match = re.findall(pattern, content)
            latitude = match[0]
            longitude = match[1]
            df_raw.loc[i, f'10_rag_{idx}_latitude'] = latitude
            df_raw.loc[i, f'10_rag_{idx}_longitude'] = longitude
        except:
            df_raw.loc[i, f'10_rag_{idx}_latitude'] = '0.0'
            df_raw.loc[i, f'10_rag_{idx}_longitude'] = '0.0'

for i in tqdm(range(df_raw.shape[0])):
    response = rag_15_df.loc[i, 'rag_response']
    response = ast.literal_eval(response)
    for idx, content in enumerate(response):
        try:
            match = re.findall(pattern, content)
            latitude = match[0]
            longitude = match[1]
            df_raw.loc[i, f'15_rag_{idx}_latitude'] = latitude
            df_raw.loc[i, f'15_rag_{idx}_longitude'] = longitude
        except:
            df_raw.loc[i, f'15_rag_{idx}_latitude'] = '0.0'
            df_raw.loc[i, f'15_rag_{idx}_longitude'] = '0.0'

df_raw.to_csv('./data/im2gps3k/im2gps3k_prediction.csv', index=False)