Update app.py

app.py

@@ -1,20 +1,41 @@
 import gradio as gr
 import torch
+import importlib.util
 from tokenizers import Tokenizer
+from huggingface_hub import hf_hub_download
 import os
-from HROM_Trainer import HROM, CONFIG, SafetyManager
 
-def load_latest_checkpoint(model, device):
-    checkpoint_dir = "checkpoints"
-    checkpoints = [f for f in os.listdir(checkpoint_dir) if f.endswith(".pt")]
-    if not checkpoints:
-        raise FileNotFoundError("No checkpoints found.")
-    checkpoints = sorted(checkpoints, key=lambda x: os.path.getmtime(os.path.join(checkpoint_dir, x)), reverse=True)
-    latest_checkpoint = os.path.join(checkpoint_dir, checkpoints[0])
-    checkpoint = torch.load(latest_checkpoint, map_location=device)
+# Download and import model components from HF Hub
+model_repo = "TimurHromek/HROM-V1"
+
+# 1. Import trainer module components
+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)
+HROM = trainer_module.HROM
+CONFIG = trainer_module.CONFIG
+SafetyManager = trainer_module.SafetyManager
+
+# 2. Load tokenizer
+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="checkpoints/HROM-V1.5.pt")
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+def load_model():
+    model = HROM().to(device)
+    checkpoint = torch.load(checkpoint_file, map_location=device)
     model.load_state_dict(checkpoint['model'])
+    model.eval()
     return model
 
+model = load_model()
+safety = SafetyManager(model, tokenizer)
+max_response_length = 200
+
 def generate_response(model, tokenizer, input_ids, safety_manager, max_length=200):
     device = next(model.parameters()).device
     generated_ids = input_ids.copy()
@@ -31,15 +52,6 @@ def generate_response(model, tokenizer, input_ids, safety_manager, max_length=20
         generated_ids.append(next_token)
     return generated_ids[len(input_ids):]
 
-# Initialize components once
-tokenizer = Tokenizer.from_file("tokenizer/hrom_tokenizer.json")
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-model = HROM().to(device)
-model = load_latest_checkpoint(model, device)
-model.eval()
-safety = SafetyManager(model, tokenizer)
-max_response_length = 200
-
 def process_message(user_input, chat_history, token_history):
     # Process user input
     user_turn = f"<user> {user_input} </s>"
@@ -80,7 +92,7 @@ def clear_history():
     return [], []
 
 with gr.Blocks() as demo:
-    gr.Markdown("# HROM Chatbot")
+    gr.Markdown("# HROM-V1 Chatbot")
     chatbot = gr.Chatbot(height=500)
     msg = gr.Textbox(label="Your Message")
     token_state = gr.State([])