-
Notifications
You must be signed in to change notification settings - Fork 0
/
train.py
38 lines (30 loc) · 1.39 KB
/
train.py
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
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.