from pymongo import MongoClient
from app.translate_loader import translate
from app.ner_loader import ner
from tqdm import tqdm
import numpy as np
import argparse
import random
import json

class InstructConverter:
def init(self):
CONNECTION_STRING = 'mongodb+srv://doadmin:7pfg65eH480k2wl3@db-mongodb-fibert-data-lake-1a57fff3.mongo.ondigitalocean.com/admin?tls=true&authSource=admin&replicaSet=db-mongodb-fibert-data-lake'
client = MongoClient(CONNECTION_STRING)
self.db = client['nutrition_database']['products']

def replace_entities(self, input_string, product_names):
translated_string = translate(input_string)
if isinstance(translated_string, list):
translated_string = translated_string[0]

ner_entities = ner(translated_string)
if len(product_names) < len(ner_entities):
return ""

for i, entity in enumerate(ner_entities):
translated_string = translated_string.replace(entity, product_names[i])

return translated_string

def count_entities(self, input_string):
translated_string = translate(input_string)
ner_entities = ner(translated_string)
return len(ner_entities)

def convert_ss(self, w=0.95, k=0.99):
'''
:param w: коэфицент неточности веса
'''
user_prompts = []
with open('./resources/static/data/food_requests.txt', 'r', encoding="utf-8") as file:
for user_prompt in file:
try:
if self.count_entities(user_prompt) == 1:
user_prompts.append(user_prompt)
except ValueError:
continue

with open('./resources/static/data/train1_instruct_ss.jsonl', 'w') as f:
for i, product in tqdm(enumerate(self.db.find())):

random_user_prompts = np.random.choice(user_prompts, 4)

for j, user_prompt in enumerate(random_user_prompts):
product_name = product['product_name']
weight_rand = round(random.random() * 150 + 50, 5)
product_name += f' {weight_rand} g'

if (i * 4 + j) % 20 == 0:
nutrient = np.random.choice(["proteins", "fats", "carbohydrates"])
nutrinent_dict = {'proteins': 'prots', 'fats': 'fats', 'carbohydrates': 'carbs'}
nutr_weight = weight_rand * product['macronutra'][nutrinent_dict[nutrient]] / 100
product_name += f" {nutr_weight} g of {nutrient}"

generated_user_prompt = self.replace_entities(
user_prompt,
[product_name, ]
)
if generated_user_prompt == user_prompt:
continue

product_json = {
"messages": [
{
"role": "user",
"content": generated_user_prompt,
},
{
"role": "assistant",
"content": json.dumps([
{
"name": product['product_name'],
"measure": "g",
"weight": {
"min": weight_rand * w,
"max": weight_rand / w,
},
"prots_100g": {
"min": product['macronutra']['prots'] * k,
"max": product['macronutra']['prots'] / k,
},
"fats_100g": {
"min": product['macronutra']['fats'] * k,
"max": product['macronutra']['fats'] / k,
},
"carbs_100g": {
"min": product['macronutra']['carbs'] * k,
"max": product['macronutra']['carbs'] / k,
},
"pols_100g": {
"min": product['macronutra']['polyols'] * k,
"max": product['macronutra']['polyols'] / k,
},
"acid_100g": {
"min": product['macronutra']['acids'] * k,
"max": product['macronutra']['acids'] / k,
},
"fiber_100g": {
"min": product['macronutra']['fiber'] * k,
"max": product['macronutra']['fiber'] / k,
},
"alco_100g": {
"min": product['macronutra']['alcohol'] * k,
"max": product['macronutra']['alcohol'] / k,
},
"kcal_100g": {
"min": product['custom_extra']['estimated_mean_energy_kcal'] * k,
"max": product['custom_extra']['estimated_mean_energy_kcal'] / k,
}
},
])
}
]
}

f.write(json.dumps(product_json))
f.write("\n")

def convert_bs(self, w=0.95, k=0.99):
'''
:param w: коэфицент неточности веса
:param k: коэфицент неточности нутры
'''

with open('./resources/static/data/food_requests.txt', 'r', encoding="utf-8") as file:
user_prompts = file.readlines()

cur_user_prompt_idx = 0
num_products = self.count_entities(user_prompts[cur_user_prompt_idx])
cur_products_count = 0
product_names = []

product_json = {
"messages": [
{
"role": "user",
},
{
"role": "assistant",
"content": []
}
]
}

with open('./static/data/train1_instruct_bs.jsonl', 'w') as f:
for i, product in tqdm(enumerate(self.db.find())):
if i > 0 and cur_products_count == num_products:
generated_user_prompt = self.replace_entities(
user_prompts[cur_user_prompt_idx % len(user_prompts)],
product_names
)

if generated_user_prompt != user_prompts[cur_user_prompt_idx % len(user_prompts)]:
if len(generated_user_prompt) > 0:
product_json['messages'][0]['content'] = generated_user_prompt
product_json['messages'][1]['content'] = json.dumps(product_json['messages'][1]['content'])

f.write(json.dumps(product_json))
f.write('\n')

cur_user_prompt_idx += 1
num_products = self.count_entities(user_prompts[cur_user_prompt_idx % len(user_prompts)])
cur_products_count = 0
product_names = []

product_json = {
"messages": [
{
"role": "user",
},
{
"role": "assistant",
"content": []
}
]
}

weight_rand = random.random() * 150 + 50
weight_min = weight_rand * w
weight_max = weight_rand / w

if i % 20 == 0:
nutrient = np.random.choice(["proteins", "fats", "carbohydrates"])
nutrient_dict = {'proteins': 'prots', 'fats': 'fats', 'carbohydrates': 'carbs'}
nutr_weight = product['macronutra'][nutrient_dict[nutrient]] * weight_rand / 100

product_names.append(
f"{product['product_name']} {nutr_weight} g of {nutrient}"
)
else:
product_names.append(product['product_name'])

product_json['messages'][1]['content'].append(
{
"name": product['product_name'],
"measure": "g",
"weight": {
"min": weight_min,
"max": weight_max,
},
"prots_100g": {
"min": product['macronutra']['prots'] * k,
"max": product['macronutra']['prots'] / k,
},
"fats_100g": {
"min": product['macronutra']['fats'] * k,
"max": product['macronutra']['fats'] / k,
},
"carbs_100g": {
"min": product['macronutra']['carbs'] * k,
"max": product['macronutra']['carbs'] / k,
},
"pols_100g": {
"min": product['macronutra']['polyols'] * k,
"max": product['macronutra']['polyols'] / k,
},
"acid_100g": {
"min": product['macronutra']['acids'] * k,
"max": product['macronutra']['acids'] / k,
},
"fiber_100g": {
"min": product['macronutra']['fiber'] * k,
"max": product['macronutra']['fiber'] / k,
},
"alco_100g": {
"min": product['macronutra']['alcohol'] * k,
"max": product['macronutra']['alcohol'] / k,
},
"kcal_100g": {
"min": min(
product['custom_extra']['energy_kcal'],
product['custom_extra']['estimated_mean_energy_kcal']
) * k,
"max": max(
product['custom_extra']['energy_kcal'],
product['custom_extra']['estimated_mean_energy_kcal']
) / k,
},
"prots": {
"min": product['macronutra']['prots'] * (weight_min / 100) * k,
"max": product['macronutra']['prots'] * (weight_max / 100) / k,
},
"fats": {
"min": product['macronutra']['fats'] * (weight_min / 100) * k,
"max": product['macronutra']['fats'] * (weight_max / 100) / k,
},
"carbs": {
"min": product['macronutra']['carbs'] * (weight_min / 100) * k,
"max": product['macronutra']['carbs'] * (weight_max / 100) / k,
},
"pols": {
"min": product['macronutra']['polyols'] * (weight_min / 100) * k,
"max": product['macronutra']['polyols'] * (weight_max / 100) / k,
},
"acid": {
"min": product['macronutra']['acids'] * (weight_min / 100) * k,
"max": product['macronutra']['acids'] * (weight_max / 100) / k,
},
"fiber": {
"min": product['macronutra']['fiber'] * (weight_min / 100) * k,
"max": product['macronutra']['fiber'] * (weight_max / 100) / k,
},
"alco": {
"min": product['macronutra']['alcohol'] * (weight_min / 100) * k,
"max": product['macronutra']['alcohol'] * (weight_max / 100) / k,
},
"kcal": {
"min": min(
product['custom_extra']['energy_kcal'],
product['custom_extra']['estimated_mean_energy_kcal']
) * (weight_min / 100) * k,
"max": max(
product['custom_extra']['energy_kcal'],
product['custom_extra']['estimated_mean_energy_kcal']
) * (weight_max / 100) / k,
}
}
)

if i % 20 == 0:
product_json['messages'][1]['content'][-1][nutrient_dict[nutrient]]['min'] = nutr_weight
product_json['messages'][1]['content'][-1][nutrient_dict[nutrient]]['max'] = nutr_weight

cur_products_count += 1

if name == "main":
parser = argparse.ArgumentParser()
parser.add_argument('convertion_type', type=str)
args = parser.parse_args()

converter = InstructConverter()
if args.convertion_type == 'ss':
converter.convert_ss()
elif args.convertion_type == 'bs':
converter.convert_bs()

Сделай жёсткий рефакторинг кода, в том числе:
функция ner() очень тяжелая, поэтому некрасиво сначала читать все варианты текста а потом по ним работать. Сделай так что бы сразу иттерировалась база данных, а внутри нее уже с помощью указателя бери строки по кругу из файла food_requests (Food_requests содержит на порядки меньше элементов поэтому нужно по кругу)
Обьедини функции ss и bs, сохранив оригинальный алгоритм, просто в этот раз передавай type как параметр
Сохрани оригинальные названия переменных и постарайся не писать новый код, а изменять текущий
При чтении food_requests можно сразу делать translate строки если она не на английском языке, но делай это через try
Не забудь про верхнеуровневый tqdm
Так же сократи код добавив функцию которая будет возвращать обьект min-max