Spaces:
Sleeping
Sleeping
File size: 3,274 Bytes
f17c7fd |
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 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 |
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, Dataset
from sklearn.model_selection import train_test_split
from transformers import BertTokenizer
# Custom Dataset Class for Text Classification
class TextDataset(Dataset):
def __init__(self, texts, labels, tokenizer, max_length=512):
self.texts = texts
self.labels = labels
self.tokenizer = tokenizer
self.max_length = max_length
def __len__(self):
return len(self.texts)
def __getitem__(self, idx):
text = self.texts[idx]
label = self.labels[idx]
# Tokenize text (we can use any tokenizer, like BERT tokenizer)
encoding = self.tokenizer(text, truncation=True, padding='max_length', max_length=self.max_length, return_tensors='pt')
input_ids = encoding['input_ids'].squeeze(0) # Remove the extra dimension
return {
'input_ids': input_ids,
'labels': torch.tensor(label, dtype=torch.long)
}
# Define your simple custom model (Feed Forward NN for classification)
class SimpleNN(nn.Module):
def __init__(self, vocab_size, hidden_size, output_size):
super(SimpleNN, self).__init__()
self.embedding = nn.Embedding(vocab_size, hidden_size)
self.fc1 = nn.Linear(hidden_size, 128)
self.fc2 = nn.Linear(128, output_size)
self.relu = nn.ReLU()
self.softmax = nn.Softmax(dim=1)
def forward(self, input_ids):
embedded = self.embedding(input_ids)
x = embedded.mean(dim=1) # Simplified pooling (averaging embeddings)
x = self.relu(self.fc1(x))
x = self.fc2(x)
return self.softmax(x)
# Example: Sample Dataset
texts = ["I love programming.", "I hate bugs.", "Python is great.", "I enjoy learning."]
labels = [1, 0, 1, 1] # For example, 1 for positive sentiment, 0 for negative
# Tokenizer (use any tokenizer - here, we're using a simple one for this example)
tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")
# Split into training and validation sets
train_texts, val_texts, train_labels, val_labels = train_test_split(texts, labels, test_size=0.2)
# Create dataset and dataloaders
train_dataset = TextDataset(train_texts, train_labels, tokenizer)
val_dataset = TextDataset(val_texts, val_labels, tokenizer)
train_loader = DataLoader(train_dataset, batch_size=2, shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=2)
# Initialize the model, optimizer, and loss function
model = SimpleNN(vocab_size=30522, hidden_size=256, output_size=2) # Output size = 2 for binary classification
optimizer = optim.Adam(model.parameters(), lr=0.001)
criterion = nn.CrossEntropyLoss()
# Training Loop
for epoch in range(3): # 3 epochs for example
model.train()
for batch in train_loader:
optimizer.zero_grad()
input_ids = batch['input_ids']
labels = batch['labels']
# Forward pass
outputs = model(input_ids)
loss = criterion(outputs, labels)
# Backward pass
loss.backward()
optimizer.step()
print(f"Epoch {epoch + 1}, Loss: {loss.item()}")
# Save the trained model
torch.save(model.state_dict(), 'custom_model.pth')
|