forked from bitosslab/linuxrustcommit
41 lines
1.5 KiB
Python
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() |