Update app.py

app.py

@@ -9,7 +9,7 @@ import os
 model_repo = "TimurHromek/HROM-V1"
 
 # 1. Import trainer module components
-trainer_file = hf_hub_download(repo_id=model_repo, filename="trainer-v1.6.py")
+trainer_file = hf_hub_download(repo_id=model_repo, filename="HROM-V1.5_Trainer.py")
 spec = importlib.util.spec_from_file_location("HROM_Trainer", trainer_file)
 trainer_module = importlib.util.module_from_spec(spec)
 spec.loader.exec_module(trainer_module)
@@ -18,11 +18,11 @@ CONFIG = trainer_module.CONFIG
 SafetyManager = trainer_module.SafetyManager
 
 # 2. Load tokenizer
-tokenizer_file = hf_hub_download(repo_id=model_repo, filename="HROM-V1.6/tokenizer/hrom_tokenizer.json")
+tokenizer_file = hf_hub_download(repo_id=model_repo, filename="tokenizer/hrom_tokenizer.json")
 tokenizer = Tokenizer.from_file(tokenizer_file)
 
 # 3. Load model checkpoint
-checkpoint_file = hf_hub_download(repo_id=model_repo, filename="HROM-V1.6/HROM-V1.6.pt")
+checkpoint_file = hf_hub_download(repo_id=model_repo, filename="HROM-V1.5_Trained-Model/HROM-V1.5.pt")
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 def load_model():
@@ -33,7 +33,7 @@ def load_model():
     return model
 
 model = load_model()
-safety = SafetyManager(model, tokenizer, device)
+safety = SafetyManager(model, tokenizer)
 max_response_length = 200
 
 def generate_response(model, tokenizer, input_ids, safety_manager, max_length=200):