Spaces:
Running
Running
File size: 1,423 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 35 36 37 38 39 |
from transformers import BertTokenizer, BertForSequenceClassification, AdamW
import torch
# Load a pre-trained BERT model and tokenizer
tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")
model = BertForSequenceClassification.from_pretrained("bert-base-uncased")
# Prepare your dataset (replace with your dataset loading code)
train_texts = ["Text of issue 1", "Text of issue 2", ...]
labels = [0, 1, ...] # 0 for non-relevant, 1 for relevant
# Tokenize and convert your dataset to tensors
input_ids = tokenizer(train_texts, truncation=True, padding=True, return_tensors="pt")
labels = torch.tensor(labels)
# Set up data loaders
dataset = torch.utils.data.TensorDataset(input_ids["input_ids"], input_ids["attention_mask"], labels)
train_loader = torch.utils.data.DataLoader(dataset, batch_size=32)
# Define optimizer and loss
optimizer = AdamW(model.parameters(), lr=1e-5)
loss_fn = torch.nn.CrossEntropyLoss()
# Fine-tune the model
model.train()
for epoch in range(3): # Replace with desired number of epochs
for batch in train_loader:
input_ids, attention_mask, labels = batch
optimizer.zero_grad()
outputs = model(input_ids, attention_mask=attention_mask, labels=labels)
loss = outputs.loss
loss.backward()
optimizer.step()
# Save the fine-tuned model
model.save_pretrained("/path/to/save/model")
# You can now use this model for semantic search.
|