Update app.py

app.py

@@ -1,43 +1,83 @@
-import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM   # or your model class
 import gradio as gr
+import torch
+import json
+from safetensors.torch import load_file as safe_load
+from huggingface_hub import hf_hub_download
+from app_classes import i3Model, ChunkTokenizer  # Make sure your classes file is importable
+
+# ------------------------------
+# Hugging Face Repo & Files
+# ------------------------------
+REPO_ID = "FlameF0X/i3-80m"  # Replace with your HF repo
+
+print("Downloading model files from Hugging Face...")
+model_file = hf_hub_download(REPO_ID, "model.safetensors")
+vocab_file = hf_hub_download(REPO_ID, "chunk_vocab_combined.json")
+config_file = hf_hub_download(REPO_ID, "config.json")
+
+# ------------------------------
+# Load Config
+# ------------------------------
+with open(config_file, "r") as f:
+    config = json.load(f)
 
-# load tokenizer & model
-model_name = "FlameF0X/i3-80m"   # replace with correct HF model path
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(model_name)
+# ------------------------------
+# Load Tokenizer
+# ------------------------------
+tokenizer = ChunkTokenizer()
+tokenizer.load(vocab_file)
+
+# ------------------------------
+# Initialize Model
+# ------------------------------
+device = "cuda" if torch.cuda.is_available() else "cpu"
+model = i3Model(vocab_size=tokenizer.vocab_size,
+                d_model=config.get("d_model", 512),
+                n_heads=config.get("n_heads", 16),
+                max_seq_len=config.get("max_seq_len", 512),
+                d_state=config.get("d_state", 32)).to(device)
+
+# Load weights
+state_dict = safe_load(model_file, device=device)
+model.load_state_dict(state_dict)
 model.eval()
-if torch.cuda.is_available():
-    model = model.cuda()
-
-def generate(prompt: str, max_new_tokens: int = 100, temperature: float = 1.0, top_k: int = None):
-    inputs = tokenizer(prompt, return_tensors="pt")
-    input_ids = inputs["input_ids"]
-    if torch.cuda.is_available():
-        input_ids = input_ids.cuda()
+
+# ------------------------------
+# Generation Function
+# ------------------------------
+def generate_text(prompt, max_tokens=100, temperature=1.0, top_k=40):
+    idx = torch.tensor([tokenizer.encode(prompt)], dtype=torch.long).to(device)
     with torch.no_grad():
-        output_ids = model.generate(
-            input_ids,
-            max_new_tokens=max_new_tokens,
-            temperature=temperature,
-            top_k=top_k,
-            do_sample=True
-        )
-    output = tokenizer.decode(output_ids[0], skip_special_tokens=True)
-    return output
-
-# Gradio interface
-iface = gr.Interface(
-    fn=generate,
-    inputs=[
-        gr.Textbox(label="Prompt", lines=2, placeholder="Enter prompt here..."),
-        gr.Slider(label="Max new tokens", minimum=1, maximum=500, step=1, value=100),
-        gr.Slider(label="Temperature", minimum=0.1, maximum=2.0, step=0.1, value=1.0),
-        gr.Slider(label="Top-k (0 = disabled)", minimum=0, maximum=200, step=1, value=40)
-    ],
-    outputs=gr.Textbox(label="Generated Text"),
-    title="i3-80m Generation Demo",
-    description="Interact with the i3 hybrid-architecture model."
-)
-
-iface.launch(server_name="0.0.0.0", server_port=7860)
+        out_idx = model.generate(idx, max_new_tokens=int(max_tokens),
+                                 temperature=float(temperature),
+                                 top_k=int(top_k))
+    return tokenizer.decode(out_idx[0].cpu())
+
+# ------------------------------
+# Gradio UI
+# ------------------------------
+with gr.Blocks() as demo:
+    gr.Markdown("## i3 Model Text Generator")
+    
+    with gr.Row():
+        prompt_input = gr.Textbox(label="Prompt", placeholder="Type your text here...", lines=3)
+        generate_btn = gr.Button("Generate")
+    
+    output_box = gr.Textbox(label="Generated Text", lines=10)
+    
+    with gr.Accordion("Dev Panel", open=False):
+        max_tokens_input = gr.Slider(10, 500, value=100, label="Max Tokens")
+        temperature_input = gr.Slider(0.1, 2.0, value=1.0, step=0.05, label="Temperature")
+        top_k_input = gr.Slider(1, tokenizer.vocab_size, value=40, step=1, label="Top-k Sampling")
+    
+    # Connect button
+    generate_btn.click(
+        generate_text,
+        inputs=[prompt_input, max_tokens_input, temperature_input, top_k_input],
+        outputs=[output_box]
+    )
+
+# ------------------------------
+# Launch App
+# ------------------------------
+demo.launch(share=True)