Update app.py

app.py

@@ -1,116 +1,215 @@
+# app.py (Gradio Client)
 import gradio as gr
-import torch
-import importlib.util
-from tokenizers import Tokenizer
+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
 
-# Download and import model components from HF Hub
+# --- 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
+# 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)
-HROM = trainer_module.HROM
-CONFIG = trainer_module.CONFIG
-SafetyManager = trainer_module.SafetyManager
+CONFIG = trainer_module.CONFIG # We need CONFIG["max_seq_len"]
 
-# 2. Load tokenizer
+# 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)
 
-# 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
+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.
     
-    # 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)
+    current_input_for_server = token_history.copy()
+    max_input_len_for_server = CONFIG["max_seq_len"] - max_response_length_config
     
-    # 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)
+    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:
-            assistant_text = tokenizer.decode(response_ids)
-            token_history.extend(response_ids)
+            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.")
     
-    chat_history.append((user_input, assistant_text))
-    return chat_history, token_history
+    # 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():
-    return [], []
+    # 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")
-    chatbot = gr.Chatbot(height=500)
-    msg = gr.Textbox(label="Your Message")
-    token_state = gr.State([])
+    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],
-        [chatbot, token_state],
-        queue=False
+        [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: "", None, msg
+        lambda: "", outputs=msg # Clear textbox
     )
     
     clear_btn = gr.Button("Clear Chat History")
     clear_btn.click(
         clear_history,
-        outputs=[chatbot, token_state],
+        outputs=[chatbot, token_state, seed_slider, temp_slider, top_k_slider],
         queue=False
     )
 
-demo.launch()
+demo.queue().launch(debug=True) # .queue() is important for streaming updates