GPTdoc / tests /test_semantic.py
YvesP's picture
initial load
65642c3
from src.tools.semantic_db import *
from src.control.control import get_long_id
texts_aldo = ["Aldo Moro [ˈaldo ˈmɔːro], né le 23 septembre 1916 à Maglie et mort assassiné le 9 mai 1978 à Rome "
"(ou ses environs), est un homme d'État italien membre de la Démocratie chrétienne (DC) Pendant la Seconde "
"Guerre mondiale, il est professeur de droit pénal.",
"Il est élu député en 1946, entre pour la première fois au "
"gouvernement en 1955 et dirige la Démocratie chrétienne entre 1959 et 1963. Il exerce deux fois les fonctions "
"de président du Conseil des ministres d'Italie (1963-1968 puis 1974-1976) et dirige par deux fois la diplomatie "
"italienne. Partisan du « compromis historique » entre les chrétiens-démocrates et les communistes, il est "
"enlevé en mars 1978 par les Brigades rouges.",
"Il est séquestré 55 jours et finalement assassiné par ses "
"geôliers. Les conditions de sa mort et l'incapacité des autorités de l'époque à le sauver restent des sujets "
"polémiques dans la classe politique et les médias italiens",
"Professeur de droit pénal à la faculté de droit de l'université de Bari à partir de 1940, Aldo Moro rejoint, "
"en 1941, la FUCI (Fédération universitaire des catholiques italiens) et en devient le président.",
"Après la "
"Seconde Guerre mondiale, il est élu à l'Assemblée constituante en 1946. Moro participe à la rédaction de la "
"nouvelle constitution. Il est ensuite réélu comme député à la Chambre des députés en 1948 pour y servir"
"jusqu'à sa mort.",
"Il est sous-secrétaire d'État aux Affaires étrangères (1948-1950), garde des sceaux, "
"ministre de la Justice (1955-1957), de l'Éducation nationale (1957-1958) et des Affaires étrangères "
"(1969-1972 et 1973-1974). Il est secrétaire de la Démocratie chrétienne de 1960 à 1963.",
"À partir de décembre "
"1963 et durant 4 ans et demi, il dirige le premier gouvernement italien intégrant des ministres socialistes.",
"Le centre gauche entre en crise à la fin des années 1960, affaiblissant le poids de Moro dans son parti. "
"Juriste renommé, il a été professeur de droit et de procédure pénale à la faculté de sciences politiques de "
"l'université de Rome « La Sapienza » de 1960 jusqu'à la mort."
]
texts_platini = ["Michel Platini, né le 21 juin 1955 à Jœuf (Meurthe-et-Moselle), est un footballeur international "
"français évoluant au poste de milieu de terrain du début des années 1970 jusqu'en 1987, avant de "
"devenir sélectionneur puis dirigeant sportif.",
"Meneur de jeu emblématique de l'équipe de France de 1976 à 1987, et en club, de l'AS Nancy-Lorraine, "
"de l'AS Saint-Étienne, puis de la Juventus de Turin, et auteur de 468 buts durant sa carrière,"
"Michel Platini est considéré comme un des meilleurs joueurs de football de l'histoire.",
"Le magazine France Football le désigne meilleur footballeur français du xxe siècle, devant Zinédine "
"Zidane et Raymond Kopa, tandis que la Juventus de Turin l'a élu meilleur « Bianconero » de tous les "
"temps.",
"Il est le premier footballeur à remporter le trophée du Ballon d'or trois fois consécutivement entre "
"1983 et 1985 (record battu par Lionel Messi remportant le trophée quatre fois de suite entre 2009 et "
"2012).",
"Il fait partie de l'équipe mondiale du xxe siècle établie par la Fédération internationale de football"
" association (FIFA) en 1998."]
query = "Quand est mort Aldo Moro?"
coll_name = get_long_id(-1)
today = "0531"
def test_long_id():
long_id = get_long_id(-1)
assert long_id[:4] == today
def test_create_collection():
coll = get_or_create_collection(coll_name, coll_name[:6])
assert coll.name == coll_name
assert coll.count() == 0
assert coll.metadata['date'] == coll_name[:6]
assert len(list_collections()) == 1
def test_add_texts():
coll = get_or_create_collection(coll_name, coll_name[:6])
add_texts_to_collection(coll_name=coll_name, texts=texts_aldo, file="aldo")
assert coll.name == coll_name
assert coll.count() == 8
assert coll.metadata['date'] == coll_name[:6]
assert len(list_collections()) == 1
def test_reset_collection():
coll = get_or_create_collection(coll_name, coll_name[:6])
reset_collection(coll.name)
assert coll.count() == 0
assert len(list_collections()) == 1
def test_delete_collection():
coll = get_or_create_collection(coll_name, coll_name[:6])
delete_collection(coll.name)
assert len(list_collections()) == 0
def test_delete_old_collections():
get_or_create_collection(coll_name, coll_name[:6])
get_or_create_collection("old_collection", '043012')
delete_old_collections()
count__ = len(list_collections())
assert count__ == 1
def test_query_collection():
coll = get_or_create_collection(coll_name, coll_name[:6])
add_texts_to_collection(coll_name=coll_name, texts=texts_aldo, file="aldo")
answers = query_collection(coll.name, query_=query, from_files=["aldo"], n_results=2)
assert len(answers['documents'][0]) == 2
answers = query_collection(coll.name, query_=query, from_files=["aldo"], n_results=3)
assert len(answers['documents'][0]) == 3
def test_query_collection_limit():
coll = get_or_create_collection(coll_name, coll_name[:6])
add_texts_to_collection(coll_name=coll_name, texts=texts_aldo, file="aldo")
add_texts_to_collection(coll_name=coll_name, texts=texts_platini, file="platini")
answers = query_collection(coll.name, query_=query, from_files=["aldo"], n_results=10)
assert len(answers['documents'][0]) == 8
def test_query_collection_no_answer():
coll = get_or_create_collection(coll_name, coll_name[:6])
answers = query_collection(coll.name, query_=query, from_files=["platini"], n_results=2)
assert len(answers['documents'][0]) == 2
def test_query_collection_no_answer():
coll = get_or_create_collection(coll_name, coll_name[:6])
answers = query_collection(coll.name, query_=query, from_files=["platini", "aldo"], n_results=2)
assert len(answers['documents'][0]) == 2
assert answers['metadatas'][0][0] == {'aldo': 1}, {'aldo': 1}