Spaces:
Running
Running
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') | |