train.py

"""Training for Veda Programming Assistant"""

import tensorflow as tf
from tensorflow import keras
import numpy as np
import os
import json

from model import VedaProgrammingLLM
from tokenizer import VedaTokenizer
from config import MODEL_DIR, VOCAB_SIZE, MAX_LENGTH, BATCH_SIZE


TRAINING_DATA = '''
<USER> Hello!
<ASSISTANT> Hello! Welcome to Veda Programming Assistant! I can help you with coding, explain concepts, or answer programming questions. What would you like to learn today?

<USER> What is Python?
<ASSISTANT> Python is a high-level, interpreted programming language created by Guido van Rossum in 1991. It is known for its simple, readable syntax.

Key features:
- Easy to learn and read
- Dynamic typing
- Large standard library
- Great for beginners and experts

Python is used for web development, data science, AI, automation, and more!

<USER> How do I print hello world?
<ASSISTANT> To print "Hello, World!" in Python, use the print() function:

<CODE>
print("Hello, World!")
<ENDCODE>

When you run this, it displays: Hello, World!

<USER> Write a function to add two numbers
<ASSISTANT> Here is a simple function to add two numbers:

<CODE>
def add(a, b):
    return a + b

# Example usage
result = add(5, 3)
print(result)  # Output: 8
<ENDCODE>

This function takes two parameters and returns their sum.

<USER> Explain what a function is
<ASSISTANT> A function is a reusable block of code that performs a specific task. Functions help organize code and avoid repetition.

Here is the basic structure:

<CODE>
def function_name(parameters):
    # Code block
    return result
<ENDCODE>

Parts of a function:
- def: keyword to define a function
- function_name: what you call your function
- parameters: input values (optional)
- return: sends back a result (optional)

Example:

<CODE>
def greet(name):
    message = f"Hello, {name}!"
    return message

print(greet("Alice"))  # Output: Hello, Alice!
<ENDCODE>

<USER> What is a list?
<ASSISTANT> A list is an ordered, mutable collection in Python that can hold items of any type.

Creating lists:

<CODE>
# Empty list
my_list = []

# List with items
numbers = [1, 2, 3, 4, 5]
fruits = ["apple", "banana", "cherry"]
mixed = [1, "hello", 3.14, True]
<ENDCODE>

Common operations:

<CODE>
fruits = ["apple", "banana"]

# Add items
fruits.append("cherry")
fruits.insert(0, "apricot")

# Access items
print(fruits[0])   # First item
print(fruits[-1])  # Last item

# Remove items
fruits.remove("banana")
fruits.pop()

# Length
print(len(fruits))
<ENDCODE>

<USER> Write a fibonacci function
<ASSISTANT> Here is a Fibonacci function:

Recursive approach:

<CODE>
def fibonacci(n):
    if n <= 0:
        return 0
    if n == 1:
        return 1
    return fibonacci(n-1) + fibonacci(n-2)

# Print first 10 numbers
for i in range(10):
    print(fibonacci(i), end=" ")
# Output: 0 1 1 2 3 5 8 13 21 34
<ENDCODE>

Iterative approach (more efficient):

<CODE>
def fibonacci_iterative(n):
    if n <= 0:
        return 0
    if n == 1:
        return 1
    
    a, b = 0, 1
    for _ in range(2, n + 1):
        a, b = b, a + b
    return b
<ENDCODE>

<USER> What is recursion?
<ASSISTANT> Recursion is when a function calls itself to solve a problem.

Key components:
1. Base case - stops the recursion
2. Recursive case - function calls itself

Example - countdown:

<CODE>
def countdown(n):
    if n <= 0:        # Base case
        print("Done!")
        return
    print(n)
    countdown(n - 1)  # Recursive call

countdown(5)
# Output: 5, 4, 3, 2, 1, Done!
<ENDCODE>

Another example - factorial:

<CODE>
def factorial(n):
    if n <= 1:        # Base case
        return 1
    return n * factorial(n - 1)  # Recursive call

print(factorial(5))  # Output: 120
<ENDCODE>

<USER> Write bubble sort
<ASSISTANT> Here is the Bubble Sort algorithm:

<CODE>
def bubble_sort(arr):
    n = len(arr)
    
    for i in range(n):
        swapped = False
        
        for j in range(0, n - i - 1):
            if arr[j] > arr[j + 1]:
                arr[j], arr[j + 1] = arr[j + 1], arr[j]
                swapped = True
        
        if not swapped:
            break
    
    return arr

# Example
numbers = [64, 34, 25, 12, 22, 11, 90]
print("Original:", numbers)
print("Sorted:", bubble_sort(numbers.copy()))
<ENDCODE>

Output:
Original: [64, 34, 25, 12, 22, 11, 90]
Sorted: [11, 12, 22, 25, 34, 64, 90]

<USER> Write binary search
<ASSISTANT> Here is the Binary Search algorithm:

<CODE>
def binary_search(arr, target):
    left = 0
    right = len(arr) - 1
    
    while left <= right:
        mid = (left + right) // 2
        
        if arr[mid] == target:
            return mid
        elif arr[mid] < target:
            left = mid + 1
        else:
            right = mid - 1
    
    return -1  # Not found

# Example
numbers = [1, 3, 5, 7, 9, 11, 13, 15]
print(binary_search(numbers, 7))   # Output: 3
print(binary_search(numbers, 6))   # Output: -1
<ENDCODE>

Important: The array must be sorted!
Time complexity: O(log n)

<USER> What is a class?
<ASSISTANT> A class is a blueprint for creating objects. It bundles data (attributes) and functions (methods) together.

<CODE>
class Dog:
    def __init__(self, name, age):
        self.name = name
        self.age = age
    
    def bark(self):
        return f"{self.name} says Woof!"
    
    def get_info(self):
        return f"{self.name} is {self.age} years old"

# Create objects
my_dog = Dog("Buddy", 3)
print(my_dog.bark())      # Buddy says Woof!
print(my_dog.get_info())  # Buddy is 3 years old
<ENDCODE>

Key concepts:
- __init__: constructor, runs when object is created
- self: refers to the current object
- Attributes: variables inside the class
- Methods: functions inside the class

<USER> How do I read a file?
<ASSISTANT> Here is how to read files in Python:

<CODE>
# Read entire file
with open("filename.txt", "r") as file:
    content = file.read()
    print(content)

# Read line by line
with open("filename.txt", "r") as file:
    for line in file:
        print(line.strip())

# Read all lines into list
with open("filename.txt", "r") as file:
    lines = file.readlines()
<ENDCODE>

With error handling:

<CODE>
def read_file_safely(filename):
    try:
        with open(filename, "r") as file:
            return file.read()
    except FileNotFoundError:
        print(f"File '{filename}' not found")
        return None

content = read_file_safely("myfile.txt")
<ENDCODE>

<USER> How do I handle errors?
<ASSISTANT> Python uses try-except blocks for error handling:

<CODE>
try:
    result = 10 / 0
except ZeroDivisionError:
    print("Cannot divide by zero!")
<ENDCODE>

Multiple exceptions:

<CODE>
try:
    number = int(input("Enter number: "))
    result = 10 / number
except ValueError:
    print("Not a valid number!")
except ZeroDivisionError:
    print("Cannot divide by zero!")
except Exception as e:
    print(f"Error: {e}")
<ENDCODE>

Complete structure:

<CODE>
try:
    result = risky_operation()
except SomeError:
    print("Handle error")
else:
    print("Success!")
finally:
    print("Always runs")
<ENDCODE>

<USER> Thanks!
<ASSISTANT> You are welcome! I am glad I could help. Feel free to ask me anything anytime! Happy coding!

<USER> What is the difference between list and tuple?
<ASSISTANT> The main difference is mutability:
- Lists are mutable (can be changed)
- Tuples are immutable (cannot be changed)

Syntax:

<CODE>
# List - square brackets
my_list = [1, 2, 3]

# Tuple - parentheses
my_tuple = (1, 2, 3)
<ENDCODE>

Mutability difference:

<CODE>
# Lists can be modified
my_list = [1, 2, 3]
my_list[0] = 100     # OK
my_list.append(4)    # OK

# Tuples cannot be modified
my_tuple = (1, 2, 3)
# my_tuple[0] = 100  # ERROR!
<ENDCODE>

When to use:
- List: data that will change
- Tuple: fixed data like coordinates

<USER> Hi
<ASSISTANT> Hello! How can I help you today? I can write code, explain concepts, or answer your programming questions!
'''


class VedaTrainer:
    """Trainer for Veda Programming Assistant"""
    
    def __init__(self, vocab_size: int = 8000, max_length: int = 512, batch_size: int = 4):
        self.vocab_size = vocab_size
        self.max_length = max_length
        self.batch_size = batch_size
        self.tokenizer = VedaTokenizer(vocab_size=vocab_size)
        self.model = None
    
    def prepare_data(self, extra_data: str = ""):
        """Prepare training data"""
        data = TRAINING_DATA
        if extra_data:
            data += "\n\n" + extra_data
        
        if os.path.exists("programming.txt"):
            with open("programming.txt", 'r', encoding='utf-8') as f:
                code_data = f.read()
            data += "\n\n" + code_data
        
        self.tokenizer.fit([data])
        
        all_tokens = self.tokenizer.encode(data)
        print(f"Total tokens: {len(all_tokens)}")
        
        sequences = []
        stride = self.max_length // 2
        
        for i in range(0, len(all_tokens) - self.max_length - 1, stride):
            seq = all_tokens[i:i + self.max_length + 1]
            if len(seq) == self.max_length + 1:
                sequences.append(seq)
        
        if len(sequences) < 10:
            stride = self.max_length // 4
            sequences = []
            for i in range(0, len(all_tokens) - self.max_length - 1, stride):
                seq = all_tokens[i:i + self.max_length + 1]
                if len(seq) == self.max_length + 1:
                    sequences.append(seq)
        
        print(f"Created {len(sequences)} training sequences")
        
        sequences = np.array(sequences)
        X = sequences[:, :-1]
        y = sequences[:, 1:]
        
        dataset = tf.data.Dataset.from_tensor_slices((X, y))
        dataset = dataset.shuffle(1000).batch(self.batch_size).prefetch(1)
        
        return dataset
    
    def build_model(self):
        """Build the model"""
        self.model = VedaProgrammingLLM(
            vocab_size=self.tokenizer.vocabulary_size,
            max_length=self.max_length,
            d_model=256,
            num_heads=8,
            num_layers=4,
            ff_dim=512
        )
        
        self.model.compile(
            optimizer=keras.optimizers.Adam(1e-4),
            loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
            metrics=['accuracy']
        )
        
        dummy = tf.zeros((1, self.max_length), dtype=tf.int32)
        self.model(dummy)
        
        return self.model
    
    def train(self, epochs: int = 15, save_path: str = None, extra_data: str = ""):
        """Train the model"""
        if save_path is None:
            save_path = MODEL_DIR
        
        dataset = self.prepare_data(extra_data)
        self.build_model()
        
        self.model.summary()
        
        os.makedirs(save_path, exist_ok=True)
        
        history = self.model.fit(dataset, epochs=epochs, verbose=1)
        
        self.model.save_weights(os.path.join(save_path, "weights.h5"))
        self.tokenizer.save(os.path.join(save_path, "tokenizer.json"))
        
        config = self.model.get_config()
        with open(os.path.join(save_path, "config.json"), 'w') as f:
            json.dump(config, f)
        
        print(f"Model saved to {save_path}")
        return history
    
    def generate_response(self, user_input: str, max_tokens: int = 200, temperature: float = 0.7) -> str:
        """Generate a response"""
        prompt = f"<USER> {user_input}\n<ASSISTANT>"
        
        tokens = self.tokenizer.encode(prompt)
        
        generated = self.model.generate(
            tokens,
            max_new_tokens=max_tokens,
            temperature=temperature,
            repetition_penalty=1.2
        )
        
        response = self.tokenizer.decode(generated)
        
        if "<ASSISTANT>" in response:
            response = response.split("<ASSISTANT>")[-1].strip()
        if "<USER>" in response:
            response = response.split("<USER>")[0].strip()
        
        return response


if __name__ == "__main__":
    trainer = VedaTrainer()
    trainer.train(epochs=20)
    
    print("\n" + "="*50)
    print("Testing:")
    print("="*50)
    
    tests = ["Hello!", "What is a function?", "Write a function to reverse a string"]
    
    for test in tests:
        print(f"\nUser: {test}")
        print(f"Assistant: {trainer.generate_response(test)}")