Spaces:
Sleeping
Sleeping
File size: 2,213 Bytes
b24d496 |
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 |
import argparse
import pandas as pd
from pathlib import Path
from sklearn.metrics import f1_score
from common.constants import DEVICE
from common.constants import DO_NORMALIZATION
from common.constants import COLUMN_LABELS_STR
from common.constants import COLUMN_TEXT
from components.faiss_vector_database import FaissVectorDatabase
from components.embedding_extraction import EmbeddingExtractor
def main(query):
global database
global model
# Здесь можно добавить любую обработку текста
cleaned_text = query.replace("\n", " ")
cleaned_text = 'query: ' + cleaned_text
query_tokens = model.query_tokenization(cleaned_text)
query_embeds = model.query_embed_extraction(query_tokens.to(DEVICE), DO_NORMALIZATION)[0]
query_embeds = query_embeds[None, :]
answer = database.search_transaction_map(query_embeds, K_NEIGHBORS)
return answer[0]
K_NEIGHBORS = 1
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument('--input_file',
type=Path,
default='../../data/csv/Карта проводок clear_3_3.pkl',
help='path to csv file.')
args = parser.parse_args()
global_model_path = 'intfloat/multilingual-e5-base'
# db_files_path = '../../data/csv/Карта проводок clear_3_3.pkl'
df = pd.read_pickle(args.input_file)
database = FaissVectorDatabase(args.input_file)
model = EmbeddingExtractor(global_model_path, DEVICE)
true, pred = [], []
for ind, row in df.iterrows():
answer = main(row[COLUMN_TEXT])
true.append(row[COLUMN_LABELS_STR])
pred.append(answer[COLUMN_LABELS_STR])
if row[COLUMN_LABELS_STR] != answer[COLUMN_LABELS_STR]:
print(f'True labels: {row[COLUMN_LABELS_STR]}-----Pred labels: {answer[COLUMN_LABELS_STR]}')
print('Название таблицы', row['TableName'])
print(row['DocName'], '---', answer['doc_name'])
print('-----------------------------------------')
print('macro', f1_score(true, pred, average='macro'))
print('micro', f1_score(true, pred, average='micro')) |