BwETAF 0.1__tar.gz

BwETAF-0.1/BwETAF/__init__.py

@@ -0,0 +1,47 @@
+import flax.serialization
+import os
+import json
+from huggingface_hub import hf_hub_download, create_repo, upload_file, login
+import os
+from .independent import debug_state
+from ._utils import convert_tree
+from .model import *
+
+
+@debug_state.trace_func
+def load_model(path,dtype= None):
+    with open(os.path.join(path, "understanding_good_stuff.json"), "r") as f:
+        data = json.load(f)
+    
+    model = ModelManager(data["num_heads"],data["attention_dim"],data["vocab_size"],data["num_blocks"],data["ff_dim"],data["dropout_rate"],dtype=dtype)
+    with open(os.path.join(path, "good_stuff.pkl"), "rb") as f:
+        model.params = convert_tree(dtype,flax.serialization.from_bytes(model.params, f.read()))
+    return model
+
+@debug_state.trace_func
+def load_hf(path,dtype= None):
+    model_repo = path
+    filenames = ["understanding_good_stuff.json","good_stuff.pkl","make_stuff_better.pkl"]
+    for i in filenames:
+        try:
+            print(hf_hub_download(repo_id=model_repo, filename=i,local_dir="Loaded_model"))
+        except:
+            print(f"No {i} found")
+    return load_model("Loaded_model",dtype)
+
+@debug_state.trace_func
+def push_model(repo_name, folder_path):
+    files_to_upload = ["good_stuff.pkl", "understanding_good_stuff.json","make_stuff_better.pkl"]
+    
+    create_repo(repo_name, exist_ok=True)  # Create repo if it doesn’t exist
+
+    for file_name in files_to_upload:
+        file_path = os.path.join(folder_path, file_name)
+        if os.path.isfile(file_path):  # Only upload if the file exists
+            upload_file(
+                path_or_fileobj=file_path,
+                path_in_repo=file_name,  # Save with the same filename
+                repo_id=repo_name,
+                repo_type="model",
+            )
+    print(f"Uploaded {files_to_upload} to {repo_name}")

BwETAF-0.1/BwETAF/_errors.py

@@ -0,0 +1,3 @@
+class IncorrectDtype(Exception):
+    def __init__(self,message):
+        super().__init__(message)

BwETAF-0.1/BwETAF/_utils.py

@@ -0,0 +1,75 @@
+import jax
+import optax
+import jax.numpy as jnp
+import numpy as np
+from ._errors import *
+import time
+from functools import partial
+
+
+def time_it(fn, *args):
+    t0 = time.time()
+    out = fn(*args)
+    t1 = time.time()
+    return out, t1 - t0
+
+
+def loss_fn(params, model,batch, rng):
+    """Computes sparse categorical cross-entropy loss with autoregressive shifting."""
+    inputs, mask, targets = batch  # Inputs: (batch, seq_len), Targets: (batch, seq_len)
+
+    logits = model.apply(params, inputs, mask,rngs={"dropout": rng})  # Forward pass
+    log_probs = jax.nn.log_softmax(logits)  # Convert logits to log probabilities
+
+    # Shift targets left: Model predicts targets[:, 1:] given inputs[:, :-1]
+    shifted_targets = targets[:, 1:]  # Remove first token
+    shifted_logits = log_probs[:, :-1, :]  # Remove last token prediction
+
+    # Get probability of correct class
+    target_probs = jnp.take_along_axis(shifted_logits, shifted_targets[..., None], axis=-1)[..., 0]
+
+    loss = -target_probs.mean()  # Negative log-likelihood loss
+
+    return loss
+
+def val_loss(params, loss_fn, model_struct, x,mask,y, key):
+    return loss_fn(params, model_struct, [x,mask,y], key)
+
+
+val_loss = jax.jit(
+    val_loss,
+    static_argnums=(1,2)
+)
+
+val_loss = jax.pmap(
+    val_loss,
+    in_axes=(None,None,None,0,0,0,None),
+    static_broadcasted_argnums=(1,2),
+    axis_name='batch'
+    )
+def BatchTrain(params, grad_fn, model_struct, x, mask, y, key, optimizer, opt_state):
+    # Yeh this too... Sry I have to hide it
+    return None
+
+BatchTrain = jax.jit(
+    BatchTrain,
+    static_argnums=(1, 2, 7)
+)
+
+BatchTrain = jax.pmap(  
+    BatchTrain,
+    static_broadcasted_argnums=(1, 2, 7),
+    in_axes=(None, None, None, 0, 0, 0, None, None, None),
+    axis_name="batch",
+    out_axes=None
+)
+
+def get_first(pytree):
+    return jax.tree_util.tree_map(lambda x: x[0], pytree)
+
+def convert_tree(dtype,pytree):
+    return jax.tree_util.tree_map(lambda x: x.astype(dtype),pytree)
+
+@partial(jax.jit, static_argnums=(0,))
+def call_model_jit(model_struct, params, x, mask, rngs):
+    return model_struct.apply(params, x, mask, rngs=rngs, training=False)

BwETAF-0.1/BwETAF/independent.py

@@ -0,0 +1,204 @@
+import numpy as np
+from concurrent.futures import ThreadPoolExecutor
+import traceback
+import jax
+import jax.numpy as jnp
+from ._errors import *
+import optax
+from ._utils import *
+import os
+import flax
+import time
+
+
+class Debugger():
+    def __init__(self,debug = False,path=None) -> None:
+        self.debug = debug
+        self.path = path
+        if path is not None:
+            self.logfile = open(path, "a")
+    
+    def turn_debugger_on(self):
+        self.debug = True
+        self.logfile = open(self.path, "a")
+
+    def logger(self,message,state="DEBUG"):
+        if self.debug:
+            if self.path is None:
+                print(f"[{state}] {message}")
+            else:
+                self.logfile.write(f"[{state}] {message}\n")
+                self.logfile.flush()
+
+    def trace_func(self,func):
+        def wrapper(*args, **kwargs):
+            self.logger(f"Calling {func.__name__} with args: {[type(i) for i in args]}, kwargs: {[type(i) for i in kwargs]}",state="FUNC_CALL")
+            start_time = time.time()
+            try:
+                out = func(*args, **kwargs)
+            except Exception as e:
+                self.logger(f"{func.__name__} Exception:{e}",state="ERROR")
+                self.logger(traceback.format_exc(), state="TRACEBACK")
+                raise
+            name = f"{args[0].__class__.__name__}.__init__" if func.__name__ == "__init__" else func.__name__
+            self.logger(f"{name} returned: {type(out)} with shape/info: {getattr(out, 'shape', 'uk')}, {getattr(out, 'dtype', 'uk')}, Time taken:{time.time()-start_time}s",state="FUNC_RETURN")
+            return out
+        return wrapper
+
+debug_state = Debugger()
+
+
+class Tokenization():
+    @debug_state.trace_func
+    def __init__(self,vocab="gpt2") -> None:
+        import tiktoken
+        self.stuff = tiktoken
+        self.vocab = vocab
+        self.enc = tiktoken.get_encoding(self.vocab)
+
+    @debug_state.trace_func
+    def tokenize(self,batch:list, workers:int, max_length:int):
+        self.enc = self.stuff.get_encoding(self.vocab)
+        enc = self.enc
+        eos_token = 50256
+        pad_token = 0
+
+        def encode_and_pad(text):
+            tokens = enc.encode(text, allowed_special={'<|endoftext|>'})[:max_length]
+            padded = np.full(max_length, pad_token, dtype=np.int32)
+            mask = np.zeros(max_length, dtype=np.int32)
+            padded[:len(tokens)] = tokens
+            mask[:len(tokens)] = 1  # Mark real tokens as 1
+            return padded, mask
+        
+        with ThreadPoolExecutor(max_workers=workers) as executor:
+            results = list(executor.map(encode_and_pad, batch))
+        
+        encoded_batch, mask = zip(*results)  # Split tokens and masks
+        encoded_batch = np.array(encoded_batch)
+        mask = np.array(mask)
+        
+        return encoded_batch, mask
+
+    @debug_state.trace_func
+    def tokenize_(self, batch: list):
+        self.enc = self.stuff.get_encoding(self.vocab)
+        enc = self.enc
+        eos_token = 50256
+
+        encoded_batch = []
+        mask = []
+
+        for text in batch:
+            tokens = enc.encode(text, allowed_special={'<|endoftext|>'})
+            encoded_batch.append(np.array(tokens, dtype=np.int32))
+            mask.append(np.ones(len(tokens), dtype=np.int32))  # Mask matches token length
+        
+        return jnp.array(encoded_batch), jnp.array(mask)
+    
+    @debug_state.trace_func
+    def decode(self,tokens):
+        return self.enc.decode(tokens)
+    
+class Flax_ds():
+    @debug_state.trace_func
+    def __init__(self,x_eq_y:bool) -> None:
+        self.x_eq_y = x_eq_y
+        self.x = None
+        self.mask = None
+        self.y = None
+        self.batch = None
+    
+    @debug_state.trace_func
+    def load_data(self,x,mask,y):
+        self.x = np.array(x)
+        self.mask = mask
+        if not self.x_eq_y:
+            self.y = np.array(y)
+    
+    @debug_state.trace_func
+    def batch_it_(self,batch_size):
+        if not self.x_eq_y:
+            self.x = jnp.array(self.x)
+            self.mask = jnp.array(self.mask)
+            self.y = jnp.array(self.y)
+            seq_len = self.x.shape[1]
+            
+            n_batches = len(self.x) // batch_size
+            num_devices = jax.device_count()
+            
+            x_batch = [self.x[i * batch_size:(i + 1) * batch_size] for i in range(n_batches)]
+            mask_batch = [self.mask[i * batch_size:(i + 1) * batch_size] for i in range(n_batches)]
+            y_batch = [self.y[i * batch_size:(i + 1) * batch_size] for i in range(n_batches)]
+            
+            
+            self.batch = [[i.reshape(num_devices,-1, seq_len), j.reshape(num_devices,-1, seq_len), k.reshape(num_devices,-1, seq_len)] for i, j, k in zip(x_batch, mask_batch, y_batch)]
+
+            del self.x, self.mask, self.y
+            return self.batch
+        
+        else:
+            self.x = jnp.array(self.x)
+            self.mask = jnp.array(self.mask)
+            seq_len = self.x.shape[1]
+            
+            n_batches = len(self.x) // batch_size
+            num_devices = jax.device_count()
+            
+            x_batch = [self.x[i * batch_size:(i + 1) * batch_size] for i in range(n_batches)]
+            mask_batch = [self.mask[i * batch_size:(i + 1) * batch_size] for i in range(n_batches)]
+            
+            
+            self.batch = [[i.reshape(num_devices,-1, seq_len), j.reshape(num_devices,-1, seq_len)] for i, j in zip(x_batch, mask_batch)]
+
+            del self.x, self.mask
+            return self.batch
+        
+    
+    def __len__(self):
+        return len(self.batch)
+
+    def stream_it(self):
+        if self.batch == None:
+            IncorrectDtype("Bruh... You forgot to run '.batch_it' before trying to stream it.... T~T")
+        if self.x_eq_y:
+            for i in self.batch:
+                yield i[0],i[1],i[0]
+        else:
+            for i in self.batch:
+                yield i[0],i[1],i[2]
+    
+    @property
+    def gimme_the_data(self):
+        return self.batch
+    
+class Optimizer():
+    @debug_state.trace_func
+    def __init__(self,optimizer,lr,lrf,batches,epochs,params):
+        decay_rate = (lrf / lr) ** (1 / (batches * epochs))
+        self.lr_schedule = optax.exponential_decay(
+            init_value=lr,
+            transition_steps=1,
+            decay_rate=decay_rate,
+            staircase=False  # Smooth decay
+        )
+        self.optimizer = optimizer(self.lr_schedule)
+        self.state = self.optimizer.init(params)
+    
+    @debug_state.trace_func
+    def load(self,path,dtype=None):
+        try:
+            with open(os.path.join(path, "make_stuff_better.pkl"), "rb") as f:
+                self.state = flax.serialization.from_bytes(self.state, f.read())
+                if dtype is not None:
+                    self.state = convert_tree(dtype,self.state)
+                print("Using loaded optimizer states")
+        except:
+            print("No optimizers states found")
+
+    @debug_state.trace_func
+    def save(self,path):
+        os.makedirs(path, exist_ok=True)
+        with open(os.path.join(path, "make_stuff_better.pkl"), "wb") as f:
+            f.write(flax.serialization.to_bytes(self.state))
+

BwETAF-0.1/BwETAF/layers.py

@@ -0,0 +1,81 @@
+import jax
+import jax.numpy as jnp
+import flax.linen as nn
+
+class PosEnc(nn.Module):
+    dim : int
+    dtype: jnp.dtype = jnp.bfloat16
+
+    @nn.compact
+    def __call__(self, x):
+        batch_size, sequence_length, _ = x.shape  
+
+        # Compute div term once (avoiding repeated `exp` calls)
+        div_term = jnp.exp(-jnp.arange(0, self.dim, 2) * (jnp.log(10000.0) / self.dim)).astype(self.dtype)
+
+        # Compute positions in one step (efficiently broadcasting)
+        position = jnp.arange(sequence_length)[:, None] * div_term  # (seq_len, emb_dim/2)
+
+        # Directly compute sine & cosine, then interleave them
+        pos_enc = jnp.zeros((sequence_length, self.dim),dtype=self.dtype)
+        pos_enc = pos_enc.at[:, 0::2].set(jnp.sin(position)).astype(self.dtype)
+        pos_enc = pos_enc.at[:, 1::2].set(jnp.cos(position)).astype(self.dtype)
+
+        # Expand for batch & return
+        return x + pos_enc[None, :, :]
+
+class Attention(nn.Module):
+    num_heads: int
+    d_model: int
+    dtype: jnp.dtype = jnp.bfloat16
+
+    def setup(self):
+        assert self.d_model % self.num_heads == 0, "d_model must be divisible by num_heads"
+        self.depth = self.d_model // self.num_heads
+        self.qkv_dense = nn.Dense(features=3 * self.d_model, kernel_init=nn.initializers.normal(stddev=0.02),dtype=self.dtype)
+        self.out_dense = nn.Dense(features=self.d_model, kernel_init=nn.initializers.normal(stddev=0.02),dtype=self.dtype)
+
+    def __call__(self, x, mask):
+        batch_size, seq_len, _ = x.shape
+        qkv = self.qkv_dense(x)  # (batch, seq_len, 3 * d_model)
+        qkv = qkv.reshape(batch_size, seq_len, 3, self.num_heads, self.depth)
+        qkv = jnp.transpose(qkv, (2, 0, 3, 1, 4))  # (3, batch, num_heads, seq_len, depth)
+        Q, K, V = qkv  # Unpacking (batch, num_heads, seq_len, depth)
+
+        # Scaled Dot-Product Attention
+        logits = jnp.einsum("bhqd,bhkd->bhqk", Q, K) / jnp.sqrt(self.depth)
+
+        if mask is not None:
+            mask = mask[:, None, :]  # Expand for broadcasting (batch, 1, seq_len, seq_len)
+            logits = jnp.where(mask, logits, -1e9)
+
+        attn_weights = jax.nn.softmax(logits, axis=-1)
+        attn_output = jnp.einsum("bhqk,bhkd->bhqd", attn_weights, V)
+
+        # Concatenate heads
+        attn_output = jnp.transpose(attn_output, (0, 2, 1, 3))  # (batch, seq_len, num_heads, depth)
+        concat_output = attn_output.reshape(batch_size, seq_len, self.d_model)  # (batch, seq_len, d_model)
+
+        return self.out_dense(concat_output)
+
+
+class Block(nn.Module):
+    num_heads : int
+    attention_dim : int
+    ff_dim : int
+    dropout_rate : float
+    dtype: jnp.dtype = jnp.bfloat16
+
+    @nn.compact
+    def __call__(self, x_inp, mask, train: bool):
+        x = nn.LayerNorm(dtype=self.dtype)(x_inp)  
+        x = Attention(self.num_heads, self.attention_dim,dtype=self.dtype)(x, mask)
+        x = nn.Dropout(self.dropout_rate)(x, deterministic=not train)
+        x_inp = x + x_inp
+
+        # Pre-LN before FFN
+        x = nn.LayerNorm(dtype=self.dtype)(x_inp)  
+        x = nn.Dense(self.ff_dim, kernel_init=nn.initializers.normal(stddev=0.02),dtype=self.dtype)(x)
+        x = nn.gelu(x)
+        x = nn.Dense(self.attention_dim, kernel_init=nn.initializers.normal(stddev=0.02),dtype=self.dtype)(x)
+        return x + x_inp

BwETAF-0.1/BwETAF/model.py

@@ -0,0 +1,186 @@
+import jax
+import jax.numpy as jnp
+import flax.linen as nn
+from .layers import *
+import json
+import os
+import flax.serialization
+import numpy as np
+from ._errors import *
+import matplotlib.pyplot as plt
+import numpy as np
+from tqdm import tqdm
+from ._utils import *
+from .independent import *
+
+
+rng = jax.random.PRNGKey(0)
+
+class Model(nn.Module):
+    num_heads: int
+    attention_dim: int
+    vocab_size: int
+    num_blocks: int
+    ff_dim: int
+    dropout_rate: float
+    dtype: jnp.dtype = jnp.float32
+
+
+    def setup(self):
+        self.emb = nn.Embed(num_embeddings=self.vocab_size,features=self.attention_dim, embedding_init=nn.initializers.normal(stddev=0.02),dtype=self.dtype)
+        self.pos_enc = PosEnc(self.attention_dim)
+        self.blocks = [Block(num_heads=self.num_heads,attention_dim=self.attention_dim,ff_dim=self.ff_dim,dropout_rate=self.dropout_rate,dtype=self.dtype)for i in range(self.num_blocks)]
+    
+    def __call__(self,x,mask,training=True):
+        mask = self.process_mask(mask)
+        x = x.astype(jnp.int32)
+        x = self.emb(x)
+        x = self.pos_enc(x)
+        for i in self.blocks:
+            x = i(x,mask,training)
+        return x @ self.emb.embedding.T
+
+
+    def process_mask(self,mask):
+        _, seq_len = mask.shape
+
+        # Create causal mask (lower triangular matrix)
+        causal_mask = jnp.tril(jnp.ones((seq_len, seq_len)))
+
+        # Reshape padding mask and apply to causal mask
+        mask = mask[:, None, :]  # (batch_size, 1, seq_len)
+        mask_sq = causal_mask[None, :, :] * mask  # (batch_size, seq_len, seq_len)
+        mask_sq = jnp.transpose(mask_sq, (0, 2, 1)) * mask
+        mask_sq = jnp.transpose(mask_sq, (0, 2, 1))
+
+        return mask_sq
+
+
+class ModelManager():
+    @debug_state.trace_func
+    def __init__(self,num_heads,attention_dim,vocab_size,num_blocks,ff_dim,dropout_rate,dtype = None) -> None:
+        self.key = jax.random.PRNGKey(0)
+        self.model_struct = Model(num_heads,attention_dim,vocab_size,num_blocks,ff_dim,dropout_rate,dtype)
+        self.params = self.model_struct.init(self.key,jax.random.normal(self.key,(2, 11)),jnp.ones((2,11))) 
+        if dtype is not None:
+            self.params = convert_tree(dtype,self.params)
+        
+        self.optimizer = None
+
+        self.data = {
+            "num_heads":num_heads,
+            "attention_dim":attention_dim,
+            "vocab_size":vocab_size,
+            "num_blocks":num_blocks,
+            "ff_dim":ff_dim,
+            "dropout_rate":dropout_rate
+        }
+
+
+    def __call__(self,input,mask):
+        return self.model_struct.apply(self.params,input,mask,rngs={"dropout": self.key_bruh},training=False)
+    
+    def jax_call(self,input,mask):
+        rngs = rngs={"dropout": self.key}
+        return call_model_jit(self.model_struct,self.params,input,mask,rngs)
+    
+
+    @property
+    def trainable_variables(self):
+        return self.params
+    
+    @property
+    def key_bruh(self):
+        self.key, subkey = jax.random.split(self.key)
+        return subkey
+    
+    @debug_state.trace_func
+    def training_setup(self,optimizer,lr,lrf,batches,epochs,state_path="",opt_state_dtype=None):
+        self.optimizer = Optimizer(optimizer,lr,lrf,batches,epochs,self.params)
+        self.optimizer.load(state_path,opt_state_dtype)
+        self.grad_fn = jax.value_and_grad(loss_fn)
+        return self.optimizer.lr_schedule
+    
+    @debug_state.trace_func
+    def train_batch(self,x,mask,y):
+        key = self.key_bruh
+        (loss, self.params, self.optimizer.state), first_time = time_it(BatchTrain,self.params,self.grad_fn,self.model_struct,x,mask,y,key,self.optimizer.optimizer,self.optimizer.state)
+        return loss , [first_time,0,0]
+    
+    @debug_state.trace_func
+    def save_model(self,name,opt_state=True):
+        os.makedirs(name, exist_ok=True)
+        with open(os.path.join(name, "good_stuff.pkl"), "wb") as f:
+            f.write(flax.serialization.to_bytes(self.trainable_variables))
+
+        with open(os.path.join(name, "understanding_good_stuff.json"),"w") as f:
+            json.dump(self.data, f, indent=2)
+        
+        if (opt_state) and (self.optimizer is not None):
+            with open(os.path.join(name, "make_stuff_better.pkl"), "wb") as f:
+                f.write(flax.serialization.to_bytes(self.optimizer.state))
+
+    @debug_state.trace_func
+    def batch_it(self, x, mask, y, batch_size, x_eq_y=True):
+        dataset = Flax_ds(x_eq_y)
+        dataset.load_data(x,mask,y)
+        dataset.batch_it_(batch_size=batch_size)
+        return dataset
+
+    @debug_state.trace_func
+    def train(self,x,mask,y,epochs,batch_size,optimizer,lr,lrf,val_x=None,val_mask=None,val_y=None,val_step=100,updates_in=1,avg_mem=1500,state_path=None):
+        pass # I don't want to expose the training stuff... So... I hope you understand :D
+    
+    
+    @debug_state.trace_func
+    def summary(self):
+        def count_params(params):
+            total = 0
+            for value in params.values():
+                if isinstance(value, dict):
+                    total += count_params(value)
+                elif hasattr(value, 'size'):
+                    total += value.size
+            return total
+        
+        for i in list(self.trainable_variables['params'].keys()):
+            print(f"{i} :{count_params(self.trainable_variables['params'].get(i, {})):,}")
+        print("-------------------")
+        print(f"Total :{count_params(self.trainable_variables['params']):,}")
+    
+    @debug_state.trace_func
+    def change_precision(self,dtype):
+        self.params = jax.tree_util.tree_map(lambda x: x.astype(dtype),self.params)
+
+    @property
+    def precision(self):
+        type_tree = jax.tree_util.tree_map(lambda x: x.dtype,self.model)
+        types = jax.tree_util.tree_leaves(type_tree)
+        if len(set(types)) == 1:
+            print(f"Model dtype:{types[0]}")
+        else:
+            print("Model contrains mixed dtypes")
+
+
+
+### Test stuff for now ok?
+### Bruh your forgot to get the better predict from googel collab ;-;
+
+@debug_state.trace_func
+def plot(losses, num_points=100, chop_off=100):
+    if chop_off >= len(losses):
+        raise ValueError("chop_off is greater than or equal to the length of losses")
+    
+    smoothed_losses = np.cumsum(losses) / (np.arange(len(losses)) + 1)
+    
+    interval = max(len(losses) // num_points, 1)
+    sampled_losses = smoothed_losses[::interval][chop_off:]
+    sampled_batches = np.arange(len(losses))[::interval][chop_off:]
+
+    # Plot
+    plt.figure(figsize=(10, 5))
+    plt.plot(sampled_batches, sampled_losses, marker='o', linestyle='-')
+    plt.xlabel('Batch')
+    plt.ylabel('Loss')
+    plt.title('Smoothed Loss over Batches')
+    plt.show()

BwETAF-0.1/BwETAF.egg-info/PKG-INFO

@@ -0,0 +1,7 @@
+Metadata-Version: 2.1
+Name: BwETAF
+Version: 0.1
+Summary: Module to load BwETAF models (Flax)
+Author: Boring._.wicked
+Classifier: Programming Language :: Python :: 3
+Classifier: License :: OSI Approved :: MIT License

BwETAF-0.1/BwETAF.egg-info/SOURCES.txt

@@ -0,0 +1,12 @@
+setup.py
+BwETAF/__init__.py
+BwETAF/_errors.py
+BwETAF/_utils.py
+BwETAF/independent.py
+BwETAF/layers.py
+BwETAF/model.py
+BwETAF.egg-info/PKG-INFO
+BwETAF.egg-info/SOURCES.txt
+BwETAF.egg-info/dependency_links.txt
+BwETAF.egg-info/requires.txt
+BwETAF.egg-info/top_level.txt

BwETAF-0.1/BwETAF.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

BwETAF-0.1/BwETAF.egg-info/requires.txt

@@ -0,0 +1,5 @@
+flax
+jax
+huggingface_hub
+optax
+numpy

BwETAF-0.1/BwETAF.egg-info/top_level.txt

@@ -0,0 +1 @@
+BwETAF

BwETAF-0.1/PKG-INFO

@@ -0,0 +1,7 @@
+Metadata-Version: 2.1
+Name: BwETAF
+Version: 0.1
+Summary: Module to load BwETAF models (Flax)
+Author: Boring._.wicked
+Classifier: Programming Language :: Python :: 3
+Classifier: License :: OSI Approved :: MIT License

BwETAF-0.1/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

BwETAF-0.1/setup.py

@@ -0,0 +1,20 @@
+from setuptools import setup, find_packages
+
+setup(
+    name="BwETAF",
+    version="0.1",
+    packages=find_packages(),
+    install_requires=[
+        "flax",
+        "jax",
+        "huggingface_hub",
+        "optax",
+        "numpy"
+    ],
+    description="Module to load BwETAF models (Flax)",
+    author="Boring._.wicked",
+    classifiers=[
+        "Programming Language :: Python :: 3",
+        "License :: OSI Approved :: MIT License",
+    ],
+)