125 lines
4.6 KiB
Python
125 lines
4.6 KiB
Python
from transformers import pipeline
|
|
from normalizer import cleaning
|
|
from elastic_helper import ElasticHelper
|
|
import transformers
|
|
import json
|
|
import datetime
|
|
import pandas as pd
|
|
from transformers import AutoTokenizer
|
|
print(transformers.__version__)
|
|
|
|
#model_checkpoint = "./BERT/findtuned_classification_model_15"# 15 epoch
|
|
model_checkpoint = '/home/gpu/tnlp/jokar/Classifier/Models/findtuned_classification_model-15'
|
|
tokenizer = AutoTokenizer.from_pretrained(model_checkpoint)
|
|
|
|
window_size = tokenizer.model_max_length#512#200
|
|
step_size = 350#100
|
|
|
|
# with open('./data/errors.txt', 'r', encoding='utf-8') as input_file:
|
|
# error_sections_id = input_file.read().splitlines()
|
|
|
|
eh_obj = ElasticHelper()
|
|
path = "/home/gpu/data_11/mj_qa_section.zip"
|
|
sections = eh_obj.iterateJsonFile(path, True)
|
|
|
|
classifier = pipeline("text-classification", model_checkpoint, framework="pt")
|
|
|
|
def get_class(sentences, top_k:int=4):
|
|
# sentences = cleaning(sentences)
|
|
out = classifier(sentences, top_k=top_k, truncation=True, max_length=window_size)
|
|
return out
|
|
|
|
def mean_classes(input_classes):
|
|
pass
|
|
all_classes = []
|
|
for cclass in input_classes:
|
|
for item in cclass:
|
|
all_classes.append({
|
|
'label': item['label'],
|
|
'score': item['score']
|
|
})
|
|
|
|
# sorted_classes = sorted(all_classes, key=lambda x: x['class'])
|
|
classes_df = pd.DataFrame(all_classes)
|
|
# گروه بندی بر اساس کلاس
|
|
grouped_df = classes_df.groupby("label").agg(
|
|
total_value=("score", "sum"), # مجموع امتیازها
|
|
count=("score", "count") # تعداد تکرار هر کلاس
|
|
).reset_index()
|
|
# تعریف فاکتور وزن بر اساس تعداد تکرار کلاس
|
|
grouped_df["weight"] = grouped_df["count"]
|
|
# بازسازی امتیاز با دخالت دادن وزن
|
|
grouped_df["score"] = grouped_df["total_value"] * grouped_df["weight"]
|
|
# حذف ستون‌های اضافی و ایجاد دیتافریم نهایی
|
|
final_df = grouped_df[["label", "count", "score"]]
|
|
# مرتب سازی دیتافریم نهایی بر اساس بالاترین امتیاز کلاسها
|
|
sorted_df = final_df.sort_values(by="score", ascending=False)
|
|
# تبدیل دیتافریم به دیکشنری
|
|
top_4_classes = sorted_df.head(4).to_dict(orient="records")
|
|
|
|
for item in top_4_classes:
|
|
# تبدیل امتیاز در مبنای درصد
|
|
item['score'] = (item['score']*100)/sorted_df['score'].sum()
|
|
item.pop('count')
|
|
|
|
return top_4_classes
|
|
|
|
def get_window_classes(text):
|
|
text_classes = []
|
|
tokens = tokenizer(text)['input_ids'][1:-1]
|
|
#print(len(tokens))
|
|
if len(tokens) > window_size:
|
|
for i in range(0, len(tokens) - window_size + 1, step_size):
|
|
start_window_slice = tokens[0: i]
|
|
window_slice = tokens[i: i + window_size]
|
|
start_char = len(tokenizer.decode(start_window_slice).replace('[UNK]', ''))
|
|
char_len = len(tokenizer.decode(window_slice).replace('[UNK]', ''))
|
|
context_slice = text[start_char: start_char + char_len]
|
|
tokens_len = len(tokenizer(context_slice)['input_ids'][1:-1])
|
|
# print(f'i: {i},token-len: {tokens_len}', flush=True)
|
|
results = get_class(context_slice)
|
|
text_classes.append(results)
|
|
|
|
text_classes = mean_classes(text_classes)
|
|
else:
|
|
text_classes = get_class(text)
|
|
|
|
return text_classes
|
|
|
|
print(f'start: {datetime.datetime.now()}')
|
|
all = 282671
|
|
new_sections_dict = {}
|
|
for index, item in enumerate(sections):
|
|
# if index > 100:
|
|
# break
|
|
id = item['id']
|
|
source = item['source']
|
|
content0 = source['content']
|
|
try:
|
|
content = cleaning(content0)
|
|
except Exception as e:
|
|
with open('./data/errors_log.txt', 'a', encoding='utf-8') as output_file:
|
|
output_file.write(id + "\n => " + str(e) + "\n\n")
|
|
continue
|
|
try:
|
|
section_classes = get_window_classes(content)
|
|
except Exception as e:
|
|
error = e
|
|
with open('./data/errors.txt', 'a', encoding='utf-8') as output_file:
|
|
output_file.write(id + "\n")
|
|
continue
|
|
# item['classes'] = section_classes
|
|
new_sections_dict[id] ={
|
|
"best-class":section_classes[0],
|
|
"other-classes": section_classes[1:]
|
|
}
|
|
print(f'section: {all}/{id}/{index+1}', flush=True)
|
|
|
|
|
|
with open('./data/all_sections_classes.json', 'w', encoding='utf-8') as output_file:
|
|
json_data = json.dumps(new_sections_dict, indent=4, ensure_ascii=False)
|
|
output_file.write(json_data)
|
|
|
|
print(f'end: {datetime.datetime.now()}')
|
|
print('finished!')
|