import argparse import math import os import tempfile import time from threading import Thread import gradio as gr import matplotlib.pyplot as plt import numpy as np import torch from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer def create_demo(model_id): hf_model_id = model_id or os.environ.get("HF_MODEL_ID", "ethicalabs/Echo-DSRN-114M") local_model = None local_tokenizer = None print(f"📦 Loading {hf_model_id}...") try: local_tokenizer = AutoTokenizer.from_pretrained(hf_model_id, trust_remote_code=True) local_model = AutoModelForCausalLM.from_pretrained( hf_model_id, trust_remote_code=True, torch_dtype=torch.float32, # CPU optimized device_map="cpu", ) local_model.eval() print("✅ Model loaded successfully on CPU.") # Metadata config = local_model.config total_params = sum(p.numel() for p in local_model.parameters()) { "Total Parameters": f"{total_params / 1e6:.1f}M", "Layers": getattr(config, "num_hidden_layers", getattr(config, "num_layers", "N/A")), "State Dim (DSRN c_t)": getattr(config, "hidden_size", 0) * getattr(config, "num_heads", 0), } except Exception as e: print(f"❌ Failed to load model: {e}") import traceback traceback.print_exc() def compress_text(text): if not text.strip() or local_model is None: return None, "Please enter some text or ensure model is loaded.", None if len(text) > 50000: return None, "❌ Error: Text exceeds the 50,000 character limit for this public demo.", None start_time = time.time() input_tokens = local_tokenizer(text, return_tensors="pt").to(local_model.device) with torch.no_grad(): outputs = local_model(**input_tokens, use_cache=True) # past_key_values is a EchoCache or tuple of states. # Echo-DSRN state is typically layer -> (h, c, k, v) or (h, c) # We want the recurrent state (c) from the LAST layer. past = outputs.past_key_values # Access the last layer's state if hasattr(past, "__getitem__"): last_layer_state = past[-1] elif hasattr(past, "states"): # EchoCache last_layer_state = past.states[-1] else: return None, "Could not extract state from model outputs.", None # index 1 is c (the slow state) c_state = last_layer_state[1] # shape (Batch, State_Dim) # Take the state of the first sequence in batch c_vector = c_state[0].cpu().numpy() elapsed = time.time() - start_time num_tokens = input_tokens.input_ids.shape[1] tps = num_tokens / elapsed # Visualize as a 2D Heatmap # State dim is e.g. 2048. We reshape to 32x64 size = c_vector.shape[0] # find closest square-like shape w = int(math.sqrt(size)) while size % w != 0 and w > 1: w -= 1 h = size // w c_matrix = c_vector.reshape((h, w)) fig, ax = plt.subplots(figsize=(6, 4)) # Use a cool-warm colormap to show positive/negative values heatmap = ax.imshow(c_matrix, cmap='coolwarm', aspect='auto') plt.colorbar(heatmap, ax=ax) ax.set_title(f"Recurrent State 'c' ({size} dims)") ax.axis('off') plt.tight_layout() stats = f"**Compression Stats:**\n- Compressed {num_tokens} tokens into {size}-dim vector.\n- Processing Speed: {tps:.1f} tokens/sec (CPU)\n- Sequence Memory Footprint: O(1)" tmp_file = tempfile.NamedTemporaryFile( delete=False, prefix="echo_dsrn_state_", suffix=".npy" ) np.save(tmp_file.name, c_vector) return fig, stats, tmp_file.name def load_file_content(file_obj): if file_obj is None: return "" try: with open(file_obj.name, "r", encoding="utf-8") as f: content = f.read() # Limit to prevent UI freezing on massive files and server OOM if len(content) > 50000: return content[:50000] + "\n\n...[TRUNCATED FOR SERVER SAFETY (50,000 char capacity)]..." return content except UnicodeDecodeError: return "❌ Error: Please upload a valid UTF-8 text file." except Exception as e: return f"❌ Error: {str(e)}" def get_text_state(text): if not text.strip() or local_model is None: return None if len(text) > 50000: raise ValueError("Text exceeds the 50,000 character limit for this public demo.") input_tokens = local_tokenizer(text, return_tensors="pt").to(local_model.device) with torch.no_grad(): outputs = local_model(**input_tokens, use_cache=True) past = outputs.past_key_values if hasattr(past, "__getitem__"): last_layer_state = past[-1] elif hasattr(past, "states"): last_layer_state = past.states[-1] else: return None c_state = last_layer_state[1] return c_state[0].cpu().numpy().flatten() def compute_similarity(file1, text1, file2, text2): try: # Resolve v1 if text1 and text1.strip(): v1 = get_text_state(text1) if v1 is None: return "❌ Failed to compute state for Text 1" elif file1 is not None: v1 = np.load(file1.name).flatten() else: return "❌ Please provide either a `.npy` file or raw text for Input 1." # Resolve v2 if text2 and text2.strip(): v2 = get_text_state(text2) if v2 is None: return "❌ Failed to compute state for Text 2" elif file2 is not None: v2 = np.load(file2.name).flatten() else: return "❌ Please provide either a `.npy` file or raw text for Input 2." if v1.shape != v2.shape: return f"❌ Shape mismatch: {v1.shape} vs {v2.shape}. They must be from the same model config." # Compute Cosine Similarity dot_product = np.dot(v1, v2) norm_v1 = np.linalg.norm(v1) norm_v2 = np.linalg.norm(v2) # Prevent division by zero if norm_v1 == 0 or norm_v2 == 0: similarity = 0.0 else: similarity = dot_product / (norm_v1 * norm_v2) # Determine color indicator based on threshold (e.g. >0.85 is high similarity) color = "#4ade80" if similarity > 0.85 else "#facc15" if similarity > 0.6 else "#f87171" html = f'''

Cosine Similarity Score

{similarity:.4f}

''' return html except Exception as e: return f"❌ Error computing similarity: {str(e)}" def generate_stream(prompt, max_tokens, temperature): if not prompt.strip() or local_model is None: yield "Please enter a prompt.", "N/A" return input_tokens = local_tokenizer(prompt, return_tensors="pt").to(local_model.device) streamer = TextIteratorStreamer(local_tokenizer, skip_prompt=True, skip_special_tokens=True) gen_kwargs = { "input_ids": input_tokens.input_ids, "max_new_tokens": max_tokens, "temperature": temperature, "do_sample": temperature > 0.0, "top_p": 0.9, "streamer": streamer, "pad_token_id": local_tokenizer.pad_token_id or 32000, } # Start generation in a separate thread thread = Thread(target=local_model.generate, kwargs=gen_kwargs) thread.start() generated_text = "" start_time = time.time() tokens_generated = 0 for new_text in streamer: generated_text += new_text tokens_generated += 1 # Estimate speed elapsed = time.time() - start_time tps = tokens_generated / elapsed if elapsed > 0 else 0 stats = f"{tps:.1f} tokens/sec on CPU" yield generated_text, stats # Final exact reading elapsed = time.time() - start_time tps = tokens_generated / elapsed if elapsed > 0 else 0 stats = ( f"**Final Speed:** {tps:.1f} tokens/sec on CPU\n**Total Tokens:** {tokens_generated}" ) yield generated_text, stats # Gradio UI Layout with gr.Blocks(title="Echo-DSRN-114M Semantic Compressor Demo (CPU)") as demo: gr.Markdown("# 🗜️ Echo-DSRN-114M Semantic Compressor Demo (CPU)") gr.Markdown( "This demo highlights the **O(1) memory footprint** and **CPU execution** of the Echo-DSRN architecture.\n\n" "⚠️ **Note for Public Space:** To prevent OOM crashes on the shared CPU tier, text inputs and document uploads are strictly limited to **50,000 characters** per request." ) with gr.Tabs(): with gr.TabItem("🧩 The Compressor"): with gr.Row(): with gr.Column(scale=1): file_upload = gr.File( label="Upload Document (.txt, .md, .py)", file_types=[".txt", ".md", ".csv", ".py"], ) compress_input = gr.Textbox( label="...or Paste a long document here", placeholder="Type or paste paragraphs of text...", lines=8, ) compress_btn = gr.Button("Compress to Fixed State", variant="primary") file_upload.change( fn=load_file_content, inputs=file_upload, outputs=compress_input ) with gr.Column(scale=1): compress_plot = gr.Plot(label="State Vector Map (c_t)") compress_stats = gr.Markdown() state_download = gr.File(label="Download State Vector (.npy)") gr.Markdown( """ **What can you do with this `.npy` file?** - **Semantic Clustering**: Load multiple state vectors in Python (`np.load()`) and cluster them (e.g., K-Means) to group similar documents. - **RAG Pre-filtering**: Use the vectors to perform Cosine Similarity searches across multiple documents. - **Cross-Attention Memory**: Treat this vector as a compressed "Gist" and feed it into a larger model (e.g., 7B) via cross-attention. - **Style Mimicry**: Train a small linear classifier on top of these vectors to detect the author's writing style or sentiment. """ ) compress_btn.click( compress_text, inputs=[compress_input], outputs=[compress_plot, compress_stats, state_download], ) with gr.TabItem("🔍 Vector Similarity"): gr.Markdown( "Upload two exported `.npy` Semantic State Vectors OR enter raw text to compute their semantic similarity." ) with gr.Row(): with gr.Column(scale=1): gr.Markdown("### Input 1") vec1_upload = gr.File( label="Upload State Vector 1 (.npy)", file_types=[".npy"] ) vec1_text = gr.Textbox( label="...or enter raw text", lines=8, value="The Apollo program, also known as Project Apollo, was the United States human spaceflight program led by NASA, which landed the first humans on the Moon in 1969. Apollo was conceived in 1960 in the Dwight D. Eisenhower presidency during Project Mercury and executed after Project Gemini. Apollo was later dedicated to President John F. Kennedy's national goal, \"before this decade is out, of landing a man on the Moon and returning him safely to the Earth\" in his address to the U.S. Congress on May 25, 1961.", ) with gr.Column(scale=1): gr.Markdown("### Input 2") vec2_upload = gr.File( label="Upload State Vector 2 (.npy)", file_types=[".npy"] ) vec2_text = gr.Textbox( label="...or enter raw text", lines=8, value="NASA's Apollo program, also referred to as Project Apollo, was the United States' human spaceflight program that successfully landed the first humans on the Moon in 1969. Apollo was conceived in 1960, during the Dwight D. Eisenhower presidency, as part of Project Mercury, and was executed after Project Gemini. Apollo was later dedicated to President John F. Kennedy's national goal of putting a man on the Moon and bringing him back home safely before the end of the 1960s. Kennedy said this in a speech to the U.S. Congress on May 25, 1961.", ) with gr.Row(): with gr.Column(scale=1): sim_btn = gr.Button("Compute Cosine Similarity", variant="primary") with gr.Column(scale=1): sim_output = gr.HTML() sim_btn.click( compute_similarity, inputs=[vec1_upload, vec1_text, vec2_upload, vec2_text], outputs=[sim_output], ) with gr.TabItem("⚡ The CPU Streamer"): with gr.Row(): with gr.Column(scale=1): stream_input = gr.Textbox( label="Prompt", lines=4, placeholder="Once upon a time..." ) with gr.Row(): max_tokens = gr.Slider(10, 512, value=128, label="Max Tokens") temperature = gr.Slider(0.0, 1.5, value=0.2, label="Temperature") stream_btn = gr.Button("Streaming Generate (CPU)", variant="primary") with gr.Column(scale=1): stream_output = gr.Textbox(label="Output", lines=8, interactive=False) stream_stats = gr.Markdown(label="Speed") stream_btn.click( generate_stream, inputs=[stream_input, max_tokens, temperature], outputs=[stream_output, stream_stats], ) return demo if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument("--model", type=str, default="ethicalabs/Echo-DSRN-114M") parser.add_argument("--port", type=int, default=7860) args = parser.parse_args() demo = create_demo(args.model) demo.queue().launch(server_port=args.port, server_name="0.0.0.0")