From 710d7661cdcda27bc5f6a95c3e0a5882f52314b2 Mon Sep 17 00:00:00 2001 From: Daniel O'Connell Date: Sat, 30 Sep 2023 18:03:08 +0200 Subject: [PATCH] handle duplicate citations --- api/src/stampy_chat/get_blocks.py | 29 ++++++++++++--------- api/tests/stampy_chat/test_get_blocks.py | 32 +++++++++++++++++------- 2 files changed, 40 insertions(+), 21 deletions(-) diff --git a/api/src/stampy_chat/get_blocks.py b/api/src/stampy_chat/get_blocks.py index 91fe738..9b59956 100644 --- a/api/src/stampy_chat/get_blocks.py +++ b/api/src/stampy_chat/get_blocks.py @@ -1,13 +1,12 @@ import dataclasses import datetime -import itertools import numpy as np import openai import regex as re import requests import time from itertools import groupby -from typing import Iterable, List, Tuple +from typing import Iterable, List from stampy_chat.env import PINECONE_NAMESPACE, REMOTE_CHAT_INSTANCE, EMBEDDING_MODEL from stampy_chat import logging @@ -54,12 +53,12 @@ def parse_block(match) -> Block: date = metadata.get('date_published') or metadata.get('date') - if isinstance(date, datetime.date): - date = date.isoformat() - elif isinstance(date, datetime.datetime): + if isinstance(date, datetime.datetime): date = date.date().isoformat() + elif isinstance(date, datetime.date): + date = date.isoformat() elif isinstance(date, (int, float)): - date = datetime.datetime.fromtimestamp(date).isoformat() + date = datetime.datetime.fromtimestamp(date).date().isoformat() authors = metadata.get('authors') if not authors and metadata.get('author'): @@ -82,16 +81,22 @@ def join_blocks(blocks: Iterable[Block]) -> List[Block]: # that the combined block has the minimum index of the blocks combined. def to_tuple(block): - return (block.id, block.title or "", block.authors or [], block.date or "", block.url or "", block.tags or "") + return (block.title or "", block.authors or [], block.date or "", block.url or "", block.tags or "") def merge_texts(blocks): return "\n.....\n".join(sorted(block.text for block in blocks)) - unified_blocks = [ - Block(*key, merge_texts(group)) - for key, group in groupby(blocks, key=to_tuple) - ] - return sorted(unified_blocks, key=to_tuple) + # There are sometimes duplicates in the dataset, but which have different ids, so the id + # is ignored when making sorting the blocks. + def make_block(key, group): + group = list(group) + # Just use the id of the first item - it doesn't matter that much in this case, as the other data points + # will be the same + return Block(group[0].id, *key, merge_texts(group)) + + blocks = sorted(blocks, key=to_tuple) + blocks = [make_block(key, group) for key, group in groupby(blocks, key=to_tuple)] + return blocks # Get the k blocks most semantically similar to the query using Pinecone. diff --git a/api/tests/stampy_chat/test_get_blocks.py b/api/tests/stampy_chat/test_get_blocks.py index 842f60e..4dbc917 100644 --- a/api/tests/stampy_chat/test_get_blocks.py +++ b/api/tests/stampy_chat/test_get_blocks.py @@ -8,22 +8,22 @@ from stampy_chat.get_blocks import Block, get_top_k_blocks, parse_block, join_bl ({}, {}), # Check dates - ({'date_published': '2023-01-02T03:04:05'}, {'date': '2023-01-02T03:04:05'}), + ({'date_published': '2023-01-01T03:04:05'}, {'date': '2023-01-01T03:04:05'}), ( {'date_published': datetime.fromisoformat('2023-01-02T03:04:05')}, - {'date': '2023-01-02T03:04:05'} - ), - ( - {'date_published': datetime.fromisoformat('2023-01-02T03:04:05').date()}, {'date': '2023-01-02'} ), ( - {'date_published': datetime.fromisoformat('2023-01-02T03:04:05').timestamp()}, - {'date': '2023-01-02T03:04:05'} + {'date_published': datetime.fromisoformat('2023-01-03T03:04:05').date()}, + {'date': '2023-01-03'} ), ( - {'date_published': int(datetime.fromisoformat('2023-01-02T03:04:05').timestamp())}, - {'date': '2023-01-02T03:04:05'} + {'date_published': datetime.fromisoformat('2023-01-04T03:04:05').timestamp()}, + {'date': '2023-01-04'} + ), + ( + {'date_published': int(datetime.fromisoformat('2023-01-05T03:04:05').timestamp())}, + {'date': '2023-01-05'} ), # Check authors @@ -84,6 +84,20 @@ def test_parse_block(match_override, block_override): Block('id3', 'title3', ['author3'], 'date3', 'url3', 'tags3', 'text3'), ] ), + ( + [ + Block('id1', 'title1', ['author1'], 'date1', 'url1', 'tags1', 'text1-1'), + Block('id3', 'title3', ['author3'], 'date3', 'url3', 'tags3', 'text3'), + Block('id1', 'title1', ['author1'], 'date1', 'url1', 'tags1', 'text1-2'), + Block('id2', 'title2', ['author2'], 'date2', 'url2', 'tags2', 'text2'), + Block('id1', 'title1', ['author1'], 'date1', 'url1', 'tags1', 'text1-3'), + ], + [ + Block('id1', 'title1', ['author1'], 'date1', 'url1', 'tags1', 'text1-1\n.....\ntext1-2\n.....\ntext1-3'), + Block('id2', 'title2', ['author2'], 'date2', 'url2', 'tags2', 'text2'), + Block('id3', 'title3', ['author3'], 'date3', 'url3', 'tags3', 'text3'), + ] + ), ]) def test_join_blocks(blocks, expected): assert list(join_blocks(blocks)) == expected