linuxrustcommit/src/train.py

41 lines
1.5 KiB
Python

import numpy as np
from sklearn.model_selection import train_test_split
from tensorflow.keras.callbacks import EarlyStopping, ModelCheckpoint
def train_model(model, X, y, validation_split=0.2, epochs=100, batch_size=32):
X_train, X_val, y_train, y_val = train_test_split(X, y, test_size=validation_split, random_state=42)
early_stopping = EarlyStopping(monitor='val_loss', patience=10, restore_best_weights=True)
model_checkpoint = ModelCheckpoint('best_model.h5', save_best_only=True, monitor='val_loss')
history = model.fit(
X_train, y_train,
validation_data=(X_val, y_val),
epochs=epochs,
batch_size=batch_size,
callbacks=[early_stopping, model_checkpoint]
)
return model, history
def evaluate_model(model, X_test, y_test):
loss, accuracy = model.evaluate(X_test, y_test)
return loss, accuracy
def main():
from model import main as model_main
model, linux_X, linux_y, rust_X, rust_y = model_main()
linux_model, linux_history = train_model(model, linux_X, linux_y)
rust_model, rust_history = train_model(model, rust_X, rust_y)
linux_loss, linux_accuracy = evaluate_model(linux_model, linux_X, linux_y)
rust_loss, rust_accuracy = evaluate_model(rust_model, rust_X, rust_y)
print(f"Linux model - Loss: {linux_loss:.4f}, Accuracy: {linux_accuracy:.4f}")
print(f"Rust model - Loss: {rust_loss:.4f}, Accuracy: {rust_accuracy:.4f}")
return linux_model, rust_model
if __name__ == '__main__':
main()