app.py

import gradio as gr
import torch
import importlib.util
from tokenizers import Tokenizer
from huggingface_hub import hf_hub_download
import os

# 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="HROM-V1.5_Trained-Model/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()
    for _ in range(max_length):
        input_tensor = torch.tensor([generated_ids], device=device)
        with torch.no_grad():
            logits = model(input_tensor)
        next_token = logits.argmax(-1)[:, -1].item()
        if next_token == tokenizer.token_to_id("</s>"):
            break
        current_text = tokenizer.decode(generated_ids + [next_token])
        if not safety_manager.content_filter(current_text):
            break
        generated_ids.append(next_token)
    return generated_ids[len(input_ids):]

def process_message(user_input, chat_history, token_history):
    # Process user input
    user_turn = f"<user> {user_input} </s>"
    user_tokens = tokenizer.encode(user_turn).ids
    token_history.extend(user_tokens)
    
    # Prepare input sequence
    input_sequence = [tokenizer.token_to_id("<s>")] + token_history
    
    # Truncate if needed
    max_input_len = CONFIG["max_seq_len"] - max_response_length
    if len(input_sequence) > max_input_len:
        input_sequence = input_sequence[-max_input_len:]
        token_history = input_sequence[1:]
    
    # Generate response
    response_ids = generate_response(model, tokenizer, input_sequence, safety, max_response_length)
    
    # Process assistant response
    assistant_text = "I couldn't generate a proper response."
    if response_ids:
        if response_ids[0] == tokenizer.token_to_id("<assistant>"):
            try:
                end_idx = response_ids.index(tokenizer.token_to_id("</s>"))
                assistant_text = tokenizer.decode(response_ids[1:end_idx])
                token_history.extend(response_ids[:end_idx+1])
            except ValueError:
                assistant_text = tokenizer.decode(response_ids[1:])
                token_history.extend(response_ids)
        else:
            assistant_text = tokenizer.decode(response_ids)
            token_history.extend(response_ids)
    
    chat_history.append((user_input, assistant_text))
    return chat_history, token_history

def clear_history():
    return [], []

# Custom CSS for styling
custom_css = """
body {
    background: linear-gradient(to bottom, #1a1a1a, #2a2a2a);
    font-family: 'Roboto', sans-serif;
    color: white;
    margin: 0;
    padding: 0;
}

.container {
    max-width: 800px;
    margin: 0 auto;
    padding: 20px;
}

.gr-chatbot {
    font-size: 16px;
    border: none;
    background-color: #1e1e1e;
    border-radius: 8px;
    padding: 10px;
}

.gr-chatbot .bubble.user {
    background-color: #2d2d2d !important;
    border-radius: 8px;
    padding: 12px;
    margin: 8px 0;
}

.gr-chatbot .bubble.assistant {
    background-color: #3d3d3d !important;
    border-radius: 8px;
    padding: 12px;
    margin: 8px 0;
}

.gr-button {
    background-color: #4CAF50;
    color: white;
    border: none;
    padding: 12px 24px;
    font-size: 16px;
    border-radius: 4px;
    cursor: pointer;
    transition: background-color 0.3s;
}

.gr-button:hover {
    background-color: #45a049;
}

.gr-text-input input {
    background-color: #2d2d2d;
    color: white;
    border: 1px solid #4CAF50;
    border-radius: 4px;
    padding: 10px;
    font-size: 16px;
}

.header {
    display: flex;
    align-items: center;
    justify-content: center;
    padding: 20px 0;
    text-align: center;
}

.header img {
    width: 60px;
    height: 60px;
    margin-right: 15px;
}

.footer {
    text-align: center;
    padding: 20px;
    font-size: 14px;
    color: #ccc;
    margin-top: 30px;
}

.title {
    font-size: 28px;
    font-weight: bold;
    color: #ffffff;
    margin: 0;
}

.subtitle {
    font-size: 16px;
    color: #cccccc;
    margin: 5px 0 0 0;
}
"""

with gr.Blocks(
    theme="dark",
    css=custom_css
) as demo:
    with gr.Column(elem_classes=["container"]):
        # Header
        with gr.Row(elem_classes=["header"]):
            gr.Image(
                value="https://huggingface.co/TimurHromek/HROM-V1/resolve/main/hrom_icon.png",
                interactive=False,
                width=60,
                height=60,
                show_label=False
            )
            with gr.Column():
                gr.Markdown("<div class='title'>HROM-V1 Chatbot</div>")
                gr.Markdown("<div class='subtitle'>Powered by Gradio and Hugging Face</div>")
        
        # Chatbot
        chatbot = gr.Chatbot(
            height=500,
            avatar_images=[
                ("https://raw.githubusercontent.com/gradio-app/gradio/main/test/test_files/test_image.png", "user"),
                ("https://raw.githubusercontent.com/gradio-app/gradio/main/test/test_files/favicon.png", "assistant")
            ],
            bubble_full_width=False
        )
        
        # Input
        msg = gr.Textbox(
            label="Your Message",
            placeholder="Type your message here...",
            lines=2
        )
        
        # Buttons
        clear_btn = gr.Button("Clear Chat History")
        
        # State
        token_state = gr.State([])
        
        # Event handlers
        msg.submit(
            process_message,
            [msg, chatbot, token_state],
            [chatbot, token_state],
            queue=False
        ).then(
            lambda: "", None, msg
        )
        
        clear_btn.click(
            clear_history,
            outputs=[chatbot, token_state],
            queue=False
        )
        
        # Footer
        gr.Markdown("<div class='footer'>© 2025 HROM-V1 | Model by Timur Hromek | Assisted by Elapt1c</div>")

if __name__ == "__main__":
    demo.launch()