Update README.md

README.md

@@ -113,6 +113,115 @@ gr_interface = gr.Interface(
 # Launch the application
 gr_interface.launch()
 
+```
+### Train Details
+
+```python
+
+# Import necessary libraries
+from datasets import load_dataset, ClassLabel
+from transformers import BertTokenizer, BertForSequenceClassification, Trainer, TrainingArguments
+import torch
+from sklearn.metrics import accuracy_score, precision_recall_fscore_support
+
+# Load dataset
+dataset = load_dataset("prithivMLmods/Spam-Text-Detect-Analysis", split="train")
+
+# Encode labels as integers
+label_mapping = {"ham": 0, "spam": 1}
+dataset = dataset.map(lambda x: {"label": label_mapping[x["Category"]]})
+dataset = dataset.rename_column("Message", "text").remove_columns(["Category"])
+
+# Convert label column to ClassLabel for stratification
+class_label = ClassLabel(names=["ham", "spam"])
+dataset = dataset.cast_column("label", class_label)
+
+# Split into train and test
+dataset = dataset.train_test_split(test_size=0.2, stratify_by_column="label")
+train_dataset = dataset["train"]
+test_dataset = dataset["test"]
+
+# Load BERT tokenizer
+tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")
+
+# Tokenize the data
+def tokenize_function(examples):
+    return tokenizer(examples["text"], padding="max_length", truncation=True, max_length=128)
+
+train_dataset = train_dataset.map(tokenize_function, batched=True)
+test_dataset = test_dataset.map(tokenize_function, batched=True)
+
+# Set format for PyTorch
+train_dataset.set_format(type="torch", columns=["input_ids", "attention_mask", "label"])
+test_dataset.set_format(type="torch", columns=["input_ids", "attention_mask", "label"])
+
+# Load pre-trained BERT model
+model = BertForSequenceClassification.from_pretrained("bert-base-uncased", num_labels=2)
+
+# Move model to GPU if available
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model.to(device)
+
+# Define evaluation metric
+def compute_metrics(eval_pred):
+    predictions, labels = eval_pred
+    predictions = torch.argmax(torch.tensor(predictions), dim=-1)
+    precision, recall, f1, _ = precision_recall_fscore_support(labels, predictions, average="binary")
+    acc = accuracy_score(labels, predictions)
+    return {"accuracy": acc, "precision": precision, "recall": recall, "f1": f1}
+
+# Training arguments
+training_args = TrainingArguments(
+    output_dir="./results",
+    evaluation_strategy="epoch",  # Evaluate after every epoch
+    save_strategy="epoch",        # Save checkpoint after every epoch
+    learning_rate=2e-5,
+    per_device_train_batch_size=16,
+    per_device_eval_batch_size=16,
+    num_train_epochs=3,
+    weight_decay=0.01,
+    logging_dir="./logs",
+    logging_steps=10,
+    load_best_model_at_end=True,
+    metric_for_best_model="accuracy",
+    greater_is_better=True
+)
+
+# Trainer
+trainer = Trainer(
+    model=model,
+    args=training_args,
+    train_dataset=train_dataset,
+    eval_dataset=test_dataset,
+    compute_metrics=compute_metrics
+)
+
+# Train the model
+trainer.train()
+
+# Evaluate the model
+results = trainer.evaluate()
+print("Evaluation Results:", results)
+
+# Save the trained model
+model.save_pretrained("./saved_model")
+tokenizer.save_pretrained("./saved_model")
+
+# Load the model for inference
+loaded_model = BertForSequenceClassification.from_pretrained("./saved_model").to(device)
+loaded_tokenizer = BertTokenizer.from_pretrained("./saved_model")
+
+# Test the model on a custom input
+def predict(text):
+    inputs = loaded_tokenizer(text, return_tensors="pt", padding="max_length", truncation=True, max_length=128)
+    inputs = {k: v.to(device) for k, v in inputs.items()}  # Move inputs to the same device as model
+    outputs = loaded_model(**inputs)
+    prediction = torch.argmax(outputs.logits, dim=-1).item()
+    return "Spam" if prediction == 1 else "Ham"
+
+# Example test
+example_text = "Congratulations! You've won a $1000 Walmart gift card. Click here to claim now."
+print("Prediction:", predict(example_text))
 ```
 
 ## **🚀 How to Train the Model**