Berufsschule/VeracityAI
Berufsschule
/
VeracityAI
Archived
3
0
Fork 0
Code Issues 9 Pull Requests Packages Projects Releases Wiki Activity

removed model error when it comes to storage

This commit is contained in:
Falko Victor Habel 2024-09-01 16:08:33 +02:00
parent 5d8ec1a01f
commit da78ad357a
1 changed files with 3 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
src/model/train.py
View File

@ -76,6 +76,8 @@ class FakeNewsModelTrainer:
total_steps = len(train_dataloader) * epochs total_steps = len(train_dataloader) * epochs
scheduler = get_linear_schedule_with_warmup(optimizer, num_warmup_steps=0, num_training_steps=total_steps) scheduler = get_linear_schedule_with_warmup(optimizer, num_warmup_steps=0, num_training_steps=total_steps)
self.model.to(self.device) # Ensure model is on the correct device
for epoch in range(epochs): for epoch in range(epochs):
self.model.train() self.model.train()
total_loss = 0 total_loss = 0
@ -128,6 +130,7 @@ class FakeNewsModelTrainer:
self.model.save_pretrained(path) self.model.save_pretrained(path)
self.tokenizer.save_pretrained(path) self.tokenizer.save_pretrained(path)
class FakeNewsInference: class FakeNewsInference:
def __init__(self, model_path): def __init__(self, model_path):
self.tokenizer = BertTokenizer.from_pretrained(model_path) self.tokenizer = BertTokenizer.from_pretrained(model_path)