Update app.py

app.py

@@ -3,59 +3,35 @@ import tensorflow as tf
 from tensorflow import keras
 from tensorflow.keras import layers
 import numpy as np
-import json
 import os
+import json
 
 # =========================================
-# 1. DATA LOADING
+# 1. SETTINGS
 # =========================================
+BLOCK_SIZE = 128
+EMBED_DIM = 256
+NUM_HEADS = 4
+FF_DIM = 512
+NUM_LAYERS = 2
+BATCH_SIZE = 32  # CPU Safe batch size
 
-# Backup text generator (multiplies text to create 200k+ chars)
-SEED_TEXT = """
-The Veda is knowledge. Knowledge is power. Wisdom is the light.
-To know the self is to know the universe. 
-Truth is one; the wise call it by many names.
-Action performed without attachment leads to liberation.
-Om Bhur Bhuva Swaha. Tat Savitur Varenyam.
-Bhargo Devasya Dhimahi. Dhiyo Yo Nah Prachodayat.
-""" * 1000 
-
-print("--- CHECKING FOR DATA ---")
-
-final_text = ""
-file_source = ""
-
-# Check if your Dad's file is uploaded
-if os.path.exists("veda.txt"):
-    print("✅ FOUND veda.txt! Loading file...")
-    with open("veda.txt", "r", encoding="utf-8", errors="ignore") as f:
-        final_text = f.read()
-    file_source = "veda.txt"
-elif os.path.exists("Veda.txt"):
-    print("✅ FOUND Veda.txt! Loading file...")
-    with open("Veda.txt", "r", encoding="utf-8", errors="ignore") as f:
-        final_text = f.read()
-    file_source = "Veda.txt"
-else:
-    print("⚠️ No file found. Using internal training data.")
-    final_text = SEED_TEXT
-    file_source = "Internal Data"
-
-print(f"Training Source: {file_source}")
-print(f"Total Characters: {len(final_text)}")
+# Paths to save the brain
+MODEL_PATH = "veda_llm.weights.h5"
+VOCAB_PATH = "vocab.json"
 
 # =========================================
-# 2. MODEL DEFINITION
+# 2. CUSTOM ARCHITECTURE (YOUR ENGINE)
 # =========================================
 @tf.keras.utils.register_keras_serializable()
-class TokenAndPositionEmbedding(tf.keras.layers.Layer):
+class TokenAndPositionEmbedding(layers.Layer):
     def __init__(self, maxlen, vocab_size, embed_dim, **kwargs):
         super().__init__(**kwargs)
         self.maxlen = maxlen
         self.vocab_size = vocab_size
         self.embed_dim = embed_dim
-        self.token_emb = tf.keras.layers.Embedding(input_dim=vocab_size, output_dim=embed_dim)
-        self.pos_emb = tf.keras.layers.Embedding(input_dim=maxlen, output_dim=embed_dim)
+        self.token_emb = layers.Embedding(input_dim=vocab_size, output_dim=embed_dim)
+        self.pos_emb = layers.Embedding(input_dim=maxlen, output_dim=embed_dim)
 
     def call(self, x):
         maxlen = tf.shape(x)[-1]
@@ -68,20 +44,17 @@ class TokenAndPositionEmbedding(tf.keras.layers.Layer):
         return config
 
 @tf.keras.utils.register_keras_serializable()
-class TransformerBlock(tf.keras.layers.Layer):
+class TransformerBlock(layers.Layer):
     def __init__(self, embed_dim, num_heads, ff_dim, rate=0.1, **kwargs):
         super().__init__(**kwargs)
         self.embed_dim = embed_dim
         self.num_heads = num_heads
         self.ff_dim = ff_dim
         self.rate = rate
-        self.att = tf.keras.layers.MultiHeadAttention(num_heads=num_heads, key_dim=embed_dim)
-        self.ffn = tf.keras.Sequential([
-            tf.keras.layers.Dense(ff_dim, activation="relu"), 
-            tf.keras.layers.Dense(embed_dim)
-        ])
-        self.ln1 = tf.keras.layers.LayerNormalization(epsilon=1e-6)
-        self.ln2 = tf.keras.layers.LayerNormalization(epsilon=1e-6)
+        self.att = layers.MultiHeadAttention(num_heads=num_heads, key_dim=embed_dim)
+        self.ffn = keras.Sequential([layers.Dense(ff_dim, activation="relu"), layers.Dense(embed_dim)])
+        self.ln1 = layers.LayerNormalization(epsilon=1e-6)
+        self.ln2 = layers.LayerNormalization(epsilon=1e-6)
 
     def call(self, inputs):
         attn_output = self.att(inputs, inputs, use_causal_mask=True)
@@ -93,100 +66,170 @@ class TransformerBlock(tf.keras.layers.Layer):
         config.update({"embed_dim": self.embed_dim, "num_heads": self.num_heads, "ff_dim": self.ff_dim, "rate": self.rate})
         return config
 
+# Function to build the model structure
+def build_llm(vocab_size):
+    inputs = layers.Input(shape=(BLOCK_SIZE,))
+    embedding_layer = TokenAndPositionEmbedding(BLOCK_SIZE, vocab_size, EMBED_DIM)
+    x = embedding_layer(inputs)
+    for _ in range(NUM_LAYERS):
+        x = TransformerBlock(EMBED_DIM, NUM_HEADS, FF_DIM)(x)
+    outputs = layers.Dense(vocab_size)(x)
+    return keras.Model(inputs=inputs, outputs=outputs)
+
+# Global Variables to hold the active brain
+current_model = None
+char2idx = {}
+idx2char = {}
+
 # =========================================
-# 3. TRAINING
+# 3. TRAINING FUNCTION (UPDATES BRAIN)
 # =========================================
-chars = sorted(list(set(final_text)))
-vocab_size = len(chars)
-char2idx = {c: i for i, c in enumerate(chars)}
-idx2char = {i: c for i, c in enumerate(chars)}
-all_ids = np.array([char2idx[c] for c in final_text])
-
-# Hyperparameters
-BATCH_SIZE = 32
-BLOCK_SIZE = 128
-EMBED_DIM = 128   
-NUM_HEADS = 4
-FF_DIM = 256
-NUM_LAYERS = 2
-EPOCHS = 3       
-
-dataset = tf.data.Dataset.from_tensor_slices(all_ids)
-dataset = dataset.batch(BLOCK_SIZE + 1, drop_remainder=True)
-dataset = dataset.map(lambda x: (x[:-1], x[1:]))
-dataset = dataset.shuffle(1000).batch(BATCH_SIZE)
-
-inputs = layers.Input(shape=(BLOCK_SIZE,))
-embedding_layer = TokenAndPositionEmbedding(BLOCK_SIZE, vocab_size, EMBED_DIM)
-x = embedding_layer(inputs)
-for _ in range(NUM_LAYERS):
-    x = TransformerBlock(EMBED_DIM, NUM_HEADS, FF_DIM)(x)
-outputs = layers.Dense(vocab_size)(x)
-model = keras.Model(inputs=inputs, outputs=outputs)
-model.compile(optimizer="adam", loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True))
-
-print(f"STARTING TRAINING...")
-try:
-    model.fit(dataset, epochs=EPOCHS)
-    print("Training Complete!")
-except Exception as e:
-    print(f"Training failed: {e}")
+def train_llm(file_obj, epochs):
+    global current_model, char2idx, idx2char
+    
+    if file_obj is None:
+        yield "Error: Please upload a .txt file first."
+        return
+
+    # 1. Read the uploaded file
+    yield f"Reading {file_obj.name}..."
+    with open(file_obj.name, 'r', encoding='utf-8', errors='ignore') as f:
+        text = f.read()
+    
+    if len(text) < BLOCK_SIZE:
+        yield "Error: Text is too short. Needs to be longer than 128 characters."
+        return
+
+    yield f"Loaded {len(text)} characters. Building Vocabulary..."
+
+    # 2. Build Vocabulary (The AI's Alphabet)
+    chars = sorted(list(set(text)))
+    vocab_size = len(chars)
+    
+    # Update global mappings
+    char2idx = {c: i for i, c in enumerate(chars)}
+    idx2char = {i: c for i, c in enumerate(chars)}
+    
+    # Save vocab immediately so Chat can use it
+    with open(VOCAB_PATH, "w") as f:
+        json.dump({"char2idx": char2idx, "idx2char": {str(k): v for k, v in idx2char.items()}}, f)
+
+    yield f"Vocab Size: {vocab_size}. Preparing Tensors..."
+
+    # 3. Create Dataset
+    all_ids = np.array([char2idx[c] for c in text])
+    text_dataset = tf.data.Dataset.from_tensor_slices(all_ids)
+    sequences = text_dataset.batch(BLOCK_SIZE + 1, drop_remainder=True)
+
+    def split_input_target(chunk):
+        return chunk[:-1], chunk[1:]
+
+    dataset = sequences.map(split_input_target).shuffle(1000).batch(BATCH_SIZE)
+
+    # 4. Initialize New Brain
+    current_model = build_llm(vocab_size)
+    optimizer = keras.optimizers.Adam(learning_rate=0.001) # High rate for fast learning
+    current_model.compile(optimizer=optimizer, loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True))
+
+    yield "Starting Training Loop..."
+
+    # 5. Training Loop
+    for epoch in range(int(epochs)):
+        history = current_model.fit(dataset, epochs=1)
+        loss = history.history['loss'][0]
+        
+        # Save Weights
+        current_model.save_weights(MODEL_PATH)
+        yield f"Epoch {epoch+1}/{epochs} Complete. Loss: {loss:.4f}"
+
+    yield "Training Complete! Go to 'Chat' tab to test your new brain."
 
 # =========================================
-# 4. CHAT GENERATION (WITH TEMPERATURE FIX)
+# 4. CHAT FUNCTION
 # =========================================
-def generate_text(prompt, length=200):
+def generate_text(prompt, length, temperature):
+    global current_model, char2idx, idx2char
+    
+    # Try to load if not in memory
+    if current_model is None:
+        if os.path.exists(MODEL_PATH) and os.path.exists(VOCAB_PATH):
+            try:
+                with open(VOCAB_PATH, "r") as f:
+                    data = json.load(f)
+                    char2idx = data["char2idx"]
+                    idx2char = {int(k): v for k, v in data["idx2char"].items()}
+                
+                vocab_size = len(char2idx)
+                current_model = build_llm(vocab_size)
+                current_model.load_weights(MODEL_PATH)
+            except:
+                return "Error: No brain found. Please go to 'Train' tab and upload a file."
+        else:
+            return "Error: Model not trained yet. Upload text in 'Train' tab."
+
     try:
+        # Pre-process prompt
         input_ids = [char2idx.get(s, 0) for s in prompt]
-        if not input_ids: return "Error: Unknown characters (not in training data)."
-            
+        if not input_ids: return "Error: Unknown characters."
+        
         input_ids = tf.convert_to_tensor([input_ids], dtype=tf.int32)
-        block_size = 128 
         result = []
 
-        # Temperature controls randomness
-        # 1.0 = Standard
-        # 0.5 = More Focused / Less Gibberish
-        # 0.2 = Very Repetitive / Safe
-        temperature = 0.5
-
         for _ in range(int(length)):
+            # Pad if prompt is short, Crop if long
             current_len = tf.shape(input_ids)[1]
-            if current_len < block_size:
-                pad_amt = block_size - current_len
+            if current_len < BLOCK_SIZE:
+                pad_amt = BLOCK_SIZE - current_len
                 padded = tf.pad(input_ids, [[0, 0], [pad_amt, 0]], constant_values=0)
             else:
-                padded = input_ids[:, -block_size:]
+                padded = input_ids[:, -BLOCK_SIZE:]
 
-            predictions = model(padded)
-            predictions = predictions[:, -1, :]
+            # Predict
+            predictions = current_model(padded)
+            predictions = predictions[:, -1, :] # Last token
             
-            # --- APPLY TEMPERATURE ---
-            # We divide logits by temperature. 
-            # Small temp (<1) makes confidence peaks higher (sharper).
+            # Apply Temperature (Creativity)
             predictions = predictions / temperature
-
+            
             predicted_id = tf.random.categorical(predictions, num_samples=1)[0, 0].numpy()
             
             input_ids = tf.concat([input_ids, [[predicted_id]]], axis=-1)
             result.append(idx2char[predicted_id])
             
         return prompt + "".join(result)
+
     except Exception as e:
         return f"Error: {str(e)}"
 
 # =========================================
 # 5. UI
 # =========================================
-iface = gr.Interface(
-    fn=generate_text,
-    inputs=[
-        gr.Textbox(label="Enter Prompt", value="The Veda is"),
-        gr.Slider(label="Length", minimum=10, maximum=500, value=200)
-    ],
-    outputs="text",
-    title="Veda AI",
-    description=f"Model trained on: {file_source} ({len(final_text)} characters)."
-)
-
-iface.launch()
+def train_wrapper(file, epochs):
+    for update in train_llm(file, epochs):
+        yield update
+
+with gr.Blocks(title="Veda LLM Trainer") as demo:
+    gr.Markdown("# Veda LLM Trainer")
+    
+    with gr.Tab("Chat"):
+        gr.Markdown("Talk to the model you trained.")
+        prompt_input = gr.Textbox(label="Prompt", value="The Veda is")
+        with gr.Row():
+            len_slider = gr.Slider(10, 500, value=200, label="Length")
+            temp_slider = gr.Slider(0.1, 2.0, value=0.6, label="Temperature (Low = Safe, High = Crazy)")
+        
+        chat_btn = gr.Button("Generate", variant="primary")
+        output_text = gr.Textbox(label="Response")
+        
+        chat_btn.click(generate_text, inputs=[prompt_input, len_slider, temp_slider], outputs=output_text)
+        
+    with gr.Tab("Train New Dataset"):
+        gr.Markdown("Upload a **.txt** file to wipe the brain and teach it new knowledge.")
+        file_input = gr.File(label="Upload Text File", file_types=[".txt"])
+        epoch_slider = gr.Slider(1, 50, value=10, step=1, label="Epochs")
+        train_btn = gr.Button("Train LLM")
+        log_box = gr.Textbox(label="Training Log")
+        
+        train_btn.click(train_wrapper, inputs=[file_input, epoch_slider], outputs=log_box)
+
+demo.launch()