from sentence_transformers import SentenceTransformer, losses, InputExample
from torch.utils.data import DataLoader
model="embeddings/all-mpnet-base-v2"
modelST = SentenceTransformer(model)


train_loss = losses.MultipleNegativesRankingLoss(model=model)

queries=["reportar un bache en mi comunidad",
                "¿Como reporto un bacheo en mi comunidad?",
                "Quiero informar sobre un hoyo en la calle cerca de mi hogar.",
                "Necesito reportar un bache en la vía cercana a mi residencia.",
                "Estoy preocupado por un hueco en la carretera que está cerca de donde vivo.",
                "Me gustaría notificar sobre un desnivel en la carretera en las inmediaciones de mi domicilio.",
                "Quiero presentar una queja sobre un socavón en la vía que conduce a mi casa.",
                "Estoy interesado en denunciar una irregularidad en la calle cercana a mi vivienda.",
                "Necesito dar aviso sobre un problema en la carretera próxima a mi residencia.",
                "Estoy observando un inconveniente en la carretera cerca de donde vivo y quisiera reportarlo.",
                "Me gustaría informar sobre un obstáculo en la vía que está cerca de mi casa.",
                "Quiero señalar un desperfecto en la carretera cercana a mi hogar.",
                "Necesito poner en conocimiento de las autoridades un bache en la calle adyacente a mi domicilio.",
                "Estoy solicitando la atención de las autoridades viales para solucionar un problema en la vía cerca de mi vivienda.",
                "Me gustaría expresar mi preocupación acerca de un bache en la carretera que está cerca de donde resido.",
                "Quiero dejar constancia de un deterioro en la calle cercana a mi casa.",
                "Estoy interesado en informar sobre un desperfecto vial en la carretera cercana a mi hogar.",
                "Necesito notificar sobre un agujero en la carretera que conduce a mi residencia.",
                "Me gustaría alertar sobre un bache en la vía cerca de mi hogar.",
                "Quiero dar aviso sobre un desnivel en la calle próxima a mi vivienda.",
                "Estoy buscando reportar un problema vial en la carretera cercana a mi casa.",
                "Necesito informar sobre un obstáculo en la vía que está cerca de mi domicilio.",
                "Me gustaría señalar una irregularidad en la carretera adyacente a mi residencia.",
                "Quiero presentar una queja sobre un desperfecto en la calle cercana a mi hogar.",
                "Estoy interesado en denunciar un socavón en la vía que conduce a mi casa.",
                "Necesito notificar sobre un inconveniente en la carretera próxima a mi residencia.",
                "Me gustaría informar sobre un deterioro en la carretera cerca de donde vivo.",
                "Quiero expresar mi preocupación acerca de un bache en la calle que está cerca de mi domicilio.",
                "Estoy solicitando la atención de las autoridades para solucionar un problema vial en las inmediaciones de mi casa.",
                "Necesito poner en conocimiento de las autoridades competentes un bache en la vía cercana a mi residencia.",
                "Me gustaría alertar sobre un agujero en la carretera cerca de mi hogar.",
                "Quiero notificar sobre un desnivel en la calle adyacente a mi domicilio.",
                "Estoy buscando reportar un desperfecto en la carretera que está cerca de mi casa.",
                "Necesito dar aviso sobre un obstáculo en la vía próxima a mi hogar.",
                "Me gustaría señalar una irregularidad en la carretera cercana a mi residencia.",
                "Quiero presentar una queja sobre un socavón en la vía que conduce a mi residencia.",
                "Estoy interesado en denunciar un inconveniente en la carretera próxima a mi casa.",
                "Necesito informar sobre un deterioro en la carretera cerca de donde vivo.",
                "Quiero expresar mi preocupación acerca de un bache en la calle que está cerca de mi domicilio.",
                "Estoy solicitando la atención de las autoridades para solucionar un problema vial en las inmediaciones de mi vivienda.",
                "Necesito poner en conocimiento de las autoridades competentes un bache en la vía cercana a mi residencia.",
                "Me gustaría alertar sobre un agujero en la carretera cerca de mi hogar.",
                "Quiero notificar sobre un desnivel en la calle adyacente a mi domicilio.",
                "Estoy buscando reportar un desperfecto en la carretera que está cerca de mi casa.",
                "Necesito dar aviso sobre un obstáculo en la vía próxima a mi hogar.",
                "Me gustaría señalar una irregularidad en la carretera cercana a mi residencia.",
                "Quiero presentar una queja sobre un socavón en la vía que conduce a mi residencia.",
                "Estoy interesado en denunciar un inconveniente en la carretera próxima a mi casa.",
                "Necesito informar sobre un deterioro en la carretera cerca de donde vivo.",
                "Quiero expresar mi preocupación acerca de un bache en la calle que está cerca de mi domicilio.",
                "Estoy solicitando la atención de las autoridades para solucionar un problema vial en las inmediaciones de mi vivienda.",
                "Necesito poner en conocimiento de las autoridades competentes un bache en la vía cercana a mi residencia."
                ]

train_examples = []
for q in queries:
    train_examples.append(InputExample(texts=[q, 'Reportar un bacheo']))

print(train_examples)
train_dataloader = DataLoader(train_examples, shuffle=True, batch_size=5)
print(train_dataloader)
num_epochs = 2

warmup_steps = int(len(train_dataloader) * num_epochs * 0.1) #10% of train data

modelST.fit(train_objectives=[(train_dataloader, train_loss)],
          epochs=num_epochs,
          warmup_steps=warmup_steps)