Spaces:
Running
Running
File size: 943 Bytes
a61e8b3 |
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 |
import pandas as pd
from sentence_transformers import SentenceTransformer, InputExample, losses
from torch.utils.data import DataLoader
from datasets import load_dataset
dataset = load_dataset("gopikrsmscs/torch-issues")
# Create InputExamples from your dataset
examples = []
for i, row in dataset['train'].iterrows():
title = row['Title']
body = row['Body']
examples.append(InputExample(texts=[title, body]))
# Load the pre-trained model
model = SentenceTransformer('sentence-transformers/all-mpnet-base-v2')
# Define a DataLoader for training
train_dataloader = DataLoader(examples, shuffle=True, batch_size=16)
# Fine-tune the model
train_loss = losses.CosineSimilarityLoss(model)
model.fit(
train_objectives=[(train_dataloader, train_loss)],
epochs=2, # You can adjust the number of training epochs
warmup_steps=100,
optimizer_params={'lr': 1e-4},
)
# Save the fine-tuned model
model.save('iSeBetter')
|