File size: 6,019 Bytes
c5e4524
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
def query_pinecone_sparse(
    dense_vec,
    sparse_vec,
    top_k,
    index,
    year,
    quarter,
    ticker,
    participant_type,
    threshold=0.25,
):
    if participant_type == "Company Speaker":
        participant = "Answer"
    else:
        participant = "Question"

    if year == "All":
        if quarter == "All":
            xc = index.query(
                vector=dense_vec,
                sparse_vector=sparse_vec,
                top_k=top_k,
                filter={
                    "Year": {
                        "$in": [
                            int("2020"),
                            int("2019"),
                            int("2018"),
                            int("2017"),
                            int("2016"),
                        ]
                    },
                    "Quarter": {"$in": ["Q1", "Q2", "Q3", "Q4"]},
                    "Ticker": {"$eq": ticker},
                    "QA_Flag": {"$eq": participant},
                },
                include_metadata=True,
            )
        else:
            xc = index.query(
                vector=dense_vec,
                sparse_vector=sparse_vec,
                top_k=top_k,
                filter={
                    "Year": {
                        "$in": [
                            int("2020"),
                            int("2019"),
                            int("2018"),
                            int("2017"),
                            int("2016"),
                        ]
                    },
                    "Quarter": {"$eq": quarter},
                    "Ticker": {"$eq": ticker},
                    "QA_Flag": {"$eq": participant},
                },
                include_metadata=True,
            )
    else:
        # search pinecone index for context passage with the answer
        xc = index.query(
            vector=dense_vec,
            sparse_vector=sparse_vec,
            top_k=top_k,
            filter={
                "Year": int(year),
                "Quarter": {"$eq": quarter},
                "Ticker": {"$eq": ticker},
                "QA_Flag": {"$eq": participant},
            },
            include_metadata=True,
        )
    # filter the context passages based on the score threshold
    filtered_matches = []
    for match in xc["matches"]:
        if match["score"] >= threshold:
            filtered_matches.append(match)
    xc["matches"] = filtered_matches
    return xc


def query_pinecone(
    dense_vec,
    top_k,
    index,
    year,
    quarter,
    ticker,
    participant_type,
    threshold=0.25,
):
    if participant_type == "Company Speaker":
        participant = "Answer"
    else:
        participant = "Question"

    if year == "All":
        if quarter == "All":
            xc = index.query(
                vector=dense_vec,
                top_k=top_k,
                filter={
                    "Year": {
                        "$in": [
                            int("2020"),
                            int("2019"),
                            int("2018"),
                            int("2017"),
                            int("2016"),
                        ]
                    },
                    "Quarter": {"$in": ["Q1", "Q2", "Q3", "Q4"]},
                    "Ticker": {"$eq": ticker},
                    "QA_Flag": {"$eq": participant},
                },
                include_metadata=True,
            )
        else:
            xc = index.query(
                vector=dense_vec,
                top_k=top_k,
                filter={
                    "Year": {
                        "$in": [
                            int("2020"),
                            int("2019"),
                            int("2018"),
                            int("2017"),
                            int("2016"),
                        ]
                    },
                    "Quarter": {"$eq": quarter},
                    "Ticker": {"$eq": ticker},
                    "QA_Flag": {"$eq": participant},
                },
                include_metadata=True,
            )
    else:
        # search pinecone index for context passage with the answer
        xc = index.query(
            vector=dense_vec,
            top_k=top_k,
            filter={
                "Year": int(year),
                "Quarter": {"$eq": quarter},
                "Ticker": {"$eq": ticker},
                "QA_Flag": {"$eq": participant},
            },
            include_metadata=True,
        )
    # filter the context passages based on the score threshold
    filtered_matches = []
    for match in xc["matches"]:
        if match["score"] >= threshold:
            filtered_matches.append(match)
    xc["matches"] = filtered_matches
    return xc


def format_query(query_results):
    # extract passage_text from Pinecone search result
    context = [
        result["metadata"]["Text"] for result in query_results["matches"]
    ]
    return context


def sentence_id_combine(data, query_results, lag=1):
    # Extract sentence IDs from query results
    ids = [
        result["metadata"]["Sentence_id"]
        for result in query_results["matches"]
    ]
    # Generate new IDs by adding a lag value to the original IDs
    new_ids = [id + i for id in ids for i in range(-lag, lag + 1)]
    # Remove duplicates and sort the new IDs
    new_ids = sorted(set(new_ids))
    # Create a list of lookup IDs by grouping the new IDs in groups of lag*2+1
    lookup_ids = [
        new_ids[i : i + (lag * 2 + 1)]
        for i in range(0, len(new_ids), lag * 2 + 1)
    ]
    # Create a list of context sentences by joining the sentences
    #  corresponding to the lookup IDs
    context_list = [
        " ".join(
            data.loc[data["Sentence_id"].isin(lookup_id), "Text"].to_list()
        )
        for lookup_id in lookup_ids
    ]
    return context_list


def text_lookup(data, sentence_ids):
    context = ". ".join(data.iloc[sentence_ids].to_list())
    return context