Create training.py

training.py

@@ -0,0 +1,152 @@
+from datasets import load_dataset
+import pandas as pd
+import fasttext
+import numpy as np
+from sklearn.model_selection import train_test_split
+from tqdm.auto import tqdm, trange
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from sklearn.metrics import mean_squared_error
+import requests
+import gzip
+import shutil
+import os
+from safetensors.torch import save_model
+import matplotlib.pyplot as plt
+
+class FasttextEmbedRegressor(nn.Module):
+    def __init__(self, input_size=300):
+        super(FasttextEmbedRegressor, self).__init__()
+        layer_1_size = 64
+        layer_2_size = 32
+        self.fc1 = nn.Linear(input_size, layer_1_size)
+        self.fc2 = nn.Linear(layer_1_size, layer_2_size)
+        self.fc3 = nn.Linear(layer_2_size, 1)
+
+    def forward(self, x):
+        x = torch.relu(self.fc1(x))
+        x = torch.relu(self.fc2(x))
+        x = self.fc3(x)
+        return x
+
+def train_regressor(X_train, X_test, y_train, y_test, train_epochs):
+
+    # Initialize the model, loss function, and optimizer
+    input_size = X_train.shape[1]
+    model = FasttextEmbedRegressor(input_size)
+    criterion = nn.MSELoss()
+    optimizer = optim.Adam(model.parameters(), lr=0.001)
+    batch_size = 32
+
+    training_metrics = []
+
+    for epoch in trange(train_epochs):
+        model.train()
+        train_losses = []
+        for step_num, i in enumerate(trange(0, X_train.shape[0], batch_size)):
+            vectors = torch.Tensor(X_train[i:i+batch_size])
+            targets = torch.Tensor(y_train[i:i+batch_size])
+            optimizer.zero_grad()
+            outputs = model(vectors).squeeze()
+            loss = criterion(outputs, targets)
+            loss.backward()
+            optimizer.step()
+            train_losses.append(float(loss))
+            if step_num % 10 == 0:
+                model.eval()
+                test_preds = model(torch.Tensor(X_test)).detach().numpy()
+                test_mse = mean_squared_error(y_test, test_preds)
+                training_metrics.append({
+                    "epoch": epoch,
+                    "step_num": step_num,
+                    "i": i,
+                    "test_mse": test_mse,
+                    "train_loss": sum(train_losses) / len(train_losses),
+                })
+                train_losses = []
+                model.train()
+
+    return pd.DataFrame(training_metrics), model
+
+def download_file(url, filename):
+    """
+    Helper method handling downloading large files from `url` to `filename`. Returns a pointer to `filename`.
+    """
+    chunkSize = 1024
+    r = requests.get(url, stream=True)
+    with open(filename, 'wb') as f:
+        pbar = tqdm( unit="B", total=int( r.headers['Content-Length'] ) )
+        for chunk in r.iter_content(chunk_size=chunkSize): 
+            if chunk: # filter out keep-alive new chunks
+                pbar.update (len(chunk))
+                f.write(chunk)
+    return filename
+
+get_filename = lambda x: f"cc.{x}.300.bin"
+
+def download_fasttext_vectors(lang_code):
+    filename = get_filename(lang_code)
+
+    if os.path.isfile(filename):
+        return None
+
+    print(f"Downloading {lang_code} vectors")
+    download_file(f"https://dl.fbaipublicfiles.com/fasttext/vectors-crawl/{filename}.gz", f"{filename}.gz")
+
+    print(f"Unzipping {lang_code} vectors")
+    with gzip.open(f"{filename}.gz", 'rb') as f_in:
+        with open(filename, 'wb') as f_out:
+            shutil.copyfileobj(f_in, f_out)
+
+    print(f"Removing zipped {lang_code} vectors")
+    os.remove(f"{filename}.gz")
+    
+    return True
+
+def create_quality_eval_model(lang_code, train_epochs=10):
+
+    download_fasttext_vectors(lang_code)
+
+    dataset = load_dataset("lightblue/text_ratings", lang_code, split="train")
+    text_list = dataset["selected_chunk"]
+    label_float = [x / 100 for x in dataset["rating_float"]]
+
+    fasttext_model = fasttext.load_model(f"cc.{lang_code}.300.bin")
+
+    embeddings = np.stack([fasttext_model.get_sentence_vector(
+        x.replace("\n", " ")
+        ) for x in tqdm(text_list)])
+
+    X_train, X_test, y_train, y_test, text_train, text_test = train_test_split(
+            embeddings,
+            label_float,
+            text_list,
+            test_size=0.2,
+            random_state=42
+    )
+
+    metrics_df, model = train_regressor(X_train, X_test, y_train, y_test, train_epochs)
+
+    test_df = pd.DataFrame({
+        "text": text_test, 
+        "gold_score": y_test, 
+        "pred_score": model(torch.Tensor(X_test)).detach().numpy().flatten()
+    })
+
+    save_model(model, f"{lang_code}.safetensors")
+
+    os.remove(get_filename(lang_code))
+
+    return metrics_df, test_df
+
+if __name__ == '__main__':
+
+    langs = ['am', 'ar', 'bg', 'bn', 'cs', 'da', 'de', 'el', 'en', 'es', 'fa', 'fi', 'fr', 'gu', 'ha', 'hi', 'hu', 'id', 'it', 'ja', 'jv', 'kn', 'ko', 'lt', 'mr', 'nl', 'no', 'yo', 'zh']
+    
+    for l in langs:
+        print(l)
+        metrics_df, test_df = create_quality_eval_model(l, train_epochs=5)
+        print(l)
+        metrics_df[["test_mse", "train_loss"]].rolling(50).mean().plot()
+        plt.show()