app.py

# app.py (Gradio Client)
import gradio as gr
import importlib.util # Still needed for CONFIG for max_seq_len
from tokenizers import Tokenizer # Still needed for local tokenization
from huggingface_hub import hf_hub_download
import os
import requests # For making HTTP requests
import json

# --- Configuration and Local Tokenizer (Still needed for UI-side processing) ---
model_repo = "TimurHromek/HROM-V1"
INFERENCE_SERVER_URL = "http://localhost:5000/generate" # CHANGE THIS to your actual https://inference.stormsurge.xyz/generate

# 1. Import trainer module components (ONLY for CONFIG if needed locally)
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)
CONFIG = trainer_module.CONFIG # We need CONFIG["max_seq_len"]

# 2. Load tokenizer (locally for encoding user input)
tokenizer_file = hf_hub_download(repo_id=model_repo, filename="tokenizer/hrom_tokenizer.json")
tokenizer = Tokenizer.from_file(tokenizer_file)

max_response_length_config = 200 # Max tokens the *server* should generate for one response

# Model and SafetyManager are NOT loaded/used on the Gradio client side for generation anymore.

def process_message(user_input, chat_history, token_history, seed, temperature, top_k):
    if not user_input.strip():
        chat_history.append((user_input, "Please provide some input."))
        return chat_history, token_history, seed, temperature, top_k # Pass back params

    # 1. Process user input and update token_history
    user_turn_text = f"<user> {user_input} </s>"
    user_tokens = tokenizer.encode(user_turn_text).ids
    token_history.extend(user_tokens)

    # 2. Add assistant marker to token_history for the server to start generating after it
    assistant_start_token = tokenizer.token_to_id("<assistant>")
    token_history.append(assistant_start_token)

    # 3. Prepare input sequence for the server (truncation)
    # The server expects the full context it needs to start generating the assistant's reply.
    # The max_response_length_config is for the *output*, so the input can be
    # max_seq_len - max_response_length_config.
    # The token_history already includes <s> from previous turns or initial state.
    
    current_input_for_server = token_history.copy()
    max_input_len_for_server = CONFIG["max_seq_len"] - max_response_length_config
    
    if len(current_input_for_server) > max_input_len_for_server:
        # If too long, truncate from the beginning, but ensure <s> is kept if present
        # More robust: find first <s> after initial part if truncating heavily.
        # Simple truncation for now:
        num_tokens_to_remove = len(current_input_for_server) - max_input_len_for_server
        # Keep <s> if it's the first token
        if current_input_for_server and current_input_for_server[0] == tokenizer.token_to_id("<s>"):
            current_input_for_server = [tokenizer.token_to_id("<s>")] + current_input_for_server[1+num_tokens_to_remove:]
        else:
            current_input_for_server = current_input_for_server[num_tokens_to_remove:]
        
        # Update token_history to reflect the truncated version sent to server
        # This is important so the client's token_history matches what the server 'sees' as context
        token_history = current_input_for_server.copy()


    # 4. Call the inference server
    payload = {
        "token_history": current_input_for_server, # This now includes <s>...<user>...</s><assistant>
        "max_response_length": max_response_length_config,
        "temperature": temperature,
        "top_k": top_k,
        "seed": seed if seed > 0 else None # Send None if seed is 0 or negative (for random)
    }

    assistant_response_text = ""
    assistant_response_token_ids = [] # Store IDs of the assistant's response
    chat_history.append((user_input, "")) # Add user message, prepare for streaming assistant

    try:
        with requests.post(INFERENCE_SERVER_URL, json=payload, stream=True, timeout=120) as r:
            r.raise_for_status() # Raise an exception for HTTP errors
            for line in r.iter_lines():
                if line:
                    decoded_line = line.decode('utf-8')
                    if decoded_line.startswith('data: '):
                        try:
                            event_data_json = decoded_line[len('data: '):]
                            event_data = json.loads(event_data_json)

                            if event_data.get("type") == "token":
                                token_text = event_data.get("text", "")
                                token_id = event_data.get("token_id")
                                assistant_response_text += token_text
                                if token_id is not None:
                                    assistant_response_token_ids.append(token_id)
                                chat_history[-1] = (user_input, assistant_response_text)
                                yield chat_history, token_history, seed, temperature, top_k # Update UI progressively
                            
                            elif event_data.get("type") == "eos":
                                # End of sentence token received
                                eos_token_id = event_data.get("token_id")
                                if eos_token_id is not None:
                                    assistant_response_token_ids.append(eos_token_id)
                                # The server should have sent </s>. We add it to token history.
                                break # Stop processing more tokens for this response
                            
                            elif event_data.get("type") == "stop":
                                reason = event_data.get("reason", "unknown reason")
                                assistant_response_text += f"\n[Generation stopped: {reason}]"
                                chat_history[-1] = (user_input, assistant_response_text)
                                yield chat_history, token_history, seed, temperature, top_k
                                break

                            elif event_data.get("type") == "stream_end":
                                # Server explicitly signals end of stream
                                break
                            
                            elif event_data.get("type") == "error":
                                err_msg = event_data.get("message", "Unknown server error")
                                assistant_response_text += f"\n[Server Error: {err_msg}]"
                                chat_history[-1] = (user_input, assistant_response_text)
                                yield chat_history, token_history, seed, temperature, top_k
                                break

                        except json.JSONDecodeError:
                            print(f"Failed to parse JSON: {decoded_line}")
                        except Exception as e:
                            print(f"Error processing stream line: {e}")
                            assistant_response_text += f"\n[Client Error: {e}]"
                            chat_history[-1] = (user_input, assistant_response_text)
                            yield chat_history, token_history, seed, temperature, top_k
                            break # Stop on error

    except requests.exceptions.RequestException as e:
        assistant_response_text = f"Error connecting to inference server: {e}"
        chat_history[-1] = (user_input, assistant_response_text)
        # No new tokens to add to token_history from assistant
        yield chat_history, token_history, seed, temperature, top_k
        return # Exit the generator

    # After stream is complete (or broken):
    # Update the main token_history with the assistant's generated tokens
    # The assistant_start_token was already added before calling the server.
    # The assistant_response_token_ids are the tokens *after* <assistant>.
    token_history.extend(assistant_response_token_ids)

    # Ensure </s> is at the end of the assistant's part in token_history if not already
    # (The server stream should ideally send eos_token_id for this)
    if not assistant_response_token_ids or assistant_response_token_ids[-1] != tokenizer.token_to_id("</s>"):
        if event_data.get("type") != "eos": # if it wasn't already an EOS event that added it
            token_history.append(tokenizer.token_to_id("</s>"))


    # Final update after generation is fully done
    if not assistant_response_text.strip(): # If nothing was generated
        chat_history[-1] = (user_input, "I couldn't generate a proper response.")
    
    # Update seed for next turn if it was used (randomize if seed was > 0)
    # If seed was <=0, it means use random, so keep it that way.
    if seed > 0:
        new_seed = seed + 1 # Or any other logic to change the seed for next turn
    else:
        new_seed = seed # Keep as random

    yield chat_history, token_history, new_seed, temperature, top_k


def clear_history():
    # Initial token_history should start with <s>
    initial_token_history = [tokenizer.token_to_id("<s>")]
    return [], initial_token_history, -1, 0.7, 50 # Cleared history, initial tokens, default seed/temp/top_k


with gr.Blocks() as demo:
    gr.Markdown("# HROM-V1 Chatbot (Remote Inference)")
    
    with gr.Row():
        with gr.Column(scale=1):
            seed_slider = gr.Slider(minimum=-1, maximum=99999, value=-1, step=1, label="Seed (-1 for random)")
            temp_slider = gr.Slider(minimum=0.1, maximum=2.0, value=0.7, step=0.05, label="Temperature")
            top_k_slider = gr.Slider(minimum=0, maximum=100, value=50, step=1, label="Top-K (0 for no Top-K)")
        with gr.Column(scale=3):
            chatbot = gr.Chatbot(height=500, label="Chat")
            msg = gr.Textbox(label="Your Message", placeholder="Type your message here...")

    # token_state stores the *entire conversation history as token IDs*
    # It should be initialized with the <s> token.
    initial_tokens = [tokenizer.token_to_id("<s>")]
    token_state = gr.State(initial_tokens)
    
    # Parameters for generation
    # seed_state = gr.State(-1) # -1 for random
    # temp_state = gr.State(0.7)
    # top_k_state = gr.State(50) # 0 to disable

    # Chain actions: submit text -> process_message (yields updates) -> clear textbox
    msg.submit(
        process_message,
        [msg, chatbot, token_state, seed_slider, temp_slider, top_k_slider],
        [chatbot, token_state, seed_slider, temp_slider, top_k_slider], # Pass params back to update state if needed
        queue=True # Enable queue for streaming
    ).then(
        lambda: "", outputs=msg # Clear textbox
    )
    
    clear_btn = gr.Button("Clear Chat History")
    clear_btn.click(
        clear_history,
        outputs=[chatbot, token_state, seed_slider, temp_slider, top_k_slider],
        queue=False
    )

demo.queue().launch(debug=True) # .queue() is important for streaming updates