Upload 7 files

README.md

@@ -1,23 +1,3 @@
----
-license: mit
-datasets:
-- lunaopenlabs/LunaAi-dataset
-language:
-- en
-metrics:
-- character
-base_model:
-- lunaopenlabs/LunaAI
-new_version: lunaopenlabs/LunaAI
-tags:
-- code
-- ai
-- luna
-- openlabs
-- open
-- source
-- text-generation-inference
----
 # Luna AI
 
 Luna AI is an open-source AI model developed by Luna OpenLabs for text classification tasks. Leveraging the BERT architecture, this model is designed to classify text into predefined categories efficiently and accurately.
@@ -115,4 +95,4 @@ Open a pull request.
 This project is licensed under the MIT License. See the LICENSE file for details.
 
 ### Contact
-For questions, suggestions, or feedback, feel free to contact the Luna OpenLabs team at [[email protected]].
+For questions, suggestions, or feedback, feel free to contact the Luna OpenLabs team at [[email protected]].

inference.py

@@ -0,0 +1,23 @@
+# inference.py
+import torch
+from transformers import BertTokenizer
+from model.luna_model import LunaAI
+
+def predict(text, model_path='./luna_ai_model'):
+    model = LunaAI(num_classes=2)
+    model.load_state_dict(torch.load(f"{model_path}/pytorch_model.bin"))
+    model.eval()
+    
+    tokenizer = BertTokenizer.from_pretrained(model_path)
+    encoding = tokenizer.encode_plus(text, return_tensors='pt')
+    input_ids, attention_mask = encoding['input_ids'], encoding['attention_mask']
+
+    with torch.no_grad():
+        output = model(input_ids, attention_mask)
+        _, prediction = torch.max(output, dim=1)
+    return prediction.item()
+
+if __name__ == "__main__":
+    sample_text = "Sample text to classify"
+    prediction = predict(sample_text)
+    print(f"Prediction: {prediction}")

requirements.txt

@@ -1,4 +1,5 @@
 torch
 transformers
 datasets
-huggingface_hub
+huggingface_hub
+scikit-learn

save_model.py

@@ -2,11 +2,11 @@
 from model.luna_model import LunaAI
 from transformers import BertTokenizer
 
-def save_model(model):
-    model.save_pretrained('./luna_ai_model')
+def save_model(model, path='./luna_ai_model'):
+    model.save_pretrained(path)
     tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
-    tokenizer.save_pretrained('./luna_ai_model')
+    tokenizer.save_pretrained(path)
 
 if __name__ == "__main__":
-    model = LunaAI()
-    save_model(model)
+    model = LunaAI(num_classes=2)  # Adjust num_classes if necessary
+    save_model(model)

training/train.py

@@ -1,16 +1,18 @@
 # training/train.py
+import os
 import pandas as pd
-from torch.utils.data import DataLoader, Dataset
-from transformers import BertTokenizer
-from model.luna_model import LunaAI
 import torch
 import torch.nn as nn
-from transformers import AdamW
+from torch.utils.data import DataLoader, Dataset
+from transformers import BertTokenizer, AdamW
+from model.luna_model import LunaAI
+from sklearn.metrics import accuracy_score, precision_recall_fscore_support
 
 class TextDataset(Dataset):
-    def __init__(self, csv_file):
+    def __init__(self, csv_file, tokenizer, max_length=128):
         self.data = pd.read_csv(csv_file)
-        self.tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
+        self.tokenizer = tokenizer
+        self.max_length = max_length
 
     def __len__(self):
         return len(self.data)
@@ -23,7 +25,7 @@ class TextDataset(Dataset):
             add_special_tokens=True,
             return_tensors='pt',
             padding='max_length',
-            max_length=128,
+            max_length=self.max_length,
             truncation=True,
         )
         return {
@@ -32,12 +34,35 @@ class TextDataset(Dataset):
             'labels': torch.tensor(label, dtype=torch.long),
         }
 
-def train_model(model, dataset):
-    dataloader = DataLoader(dataset, batch_size=16, shuffle=True)
-    optimizer = AdamW(model.parameters(), lr=5e-5)
+def evaluate_model(model, dataloader):
+    model.eval()
+    predictions, true_labels = [], []
+    with torch.no_grad():
+        for batch in dataloader:
+            outputs = model(batch['input_ids'], batch['attention_mask'])
+            _, preds = torch.max(outputs, dim=1)
+            predictions.extend(preds.cpu().numpy())
+            true_labels.extend(batch['labels'].cpu().numpy())
 
+    accuracy = accuracy_score(true_labels, predictions)
+    precision, recall, f1, _ = precision_recall_fscore_support(true_labels, predictions, average='weighted')
+    return accuracy, precision, recall, f1
+
+def save_checkpoint(epoch, model, optimizer, loss, path="./checkpoints"):
+    os.makedirs(path, exist_ok=True)
+    torch.save({
+        'epoch': epoch,
+        'model_state_dict': model.state_dict(),
+        'optimizer_state_dict': optimizer.state_dict(),
+        'loss': loss,
+    }, os.path.join(path, f"checkpoint_epoch_{epoch}.pth"))
+
+def train_model(model, dataset, epochs=3, batch_size=16, learning_rate=5e-5):
+    dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)
+    optimizer = AdamW(model.parameters(), lr=learning_rate)
     model.train()
-    for epoch in range(3):  # Adjust the number of epochs
+    
+    for epoch in range(epochs):
         for batch in dataloader:
             input_ids = batch['input_ids']
             attention_mask = batch['attention_mask']
@@ -50,7 +75,14 @@ def train_model(model, dataset):
             optimizer.step()
             print(f'Epoch {epoch}, Loss: {loss.item()}')
 
+        save_checkpoint(epoch, model, optimizer, loss.item())
+
+        # Optional: Evaluate the model at each epoch end
+        accuracy, precision, recall, f1 = evaluate_model(model, dataloader)
+        print(f'Epoch {epoch} - Accuracy: {accuracy}, Precision: {precision}, Recall: {recall}, F1 Score: {f1}')
+
 if __name__ == "__main__":
-    dataset = TextDataset('data/dataset.csv')
-    model = LunaAI()
+    tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
+    dataset = TextDataset('data/dataset.csv', tokenizer)
+    model = LunaAI(num_classes=2)  # Adjust num_classes if necessary
     train_model(model, dataset)