Paper - Pytorch
Project description
Chai-1
An free and open source community implementation of Chai-1 in PyTorch. Paper is here
Join our discord to help us implement this paper!
Installation
pip3 install chai-one
Usage
######### example.py
import torch
from loguru import logger
from chai_one.model import ChaiOne
# Set up model parameters
dim_single = 128
dim_pairwise = 128
dim_msa = 128
dim_msa_input = 134 # Adjusted to match the expected input dimension
dim_additional_msa_feats = 2
window_size = 25
# Initialize the model
logger.info("Initializing ChaiOne model")
model = ChaiOne(
dim_single=dim_single,
dim_pairwise=dim_pairwise,
msa_depth=4,
dim_msa=dim_msa,
dim_msa_input=dim_msa_input, # Set to 134
dim_additional_msa_feats=0,
msa_pwa_heads=8,
msa_pwa_dim_head=32,
layerscale_output=False,
heads=8,
window_size=window_size,
num_memory_kv=0,
attn_layers=48,
)
# Create dummy input tensors
batch_size = 1
seq_length = 100
num_msa = 4
logger.info(
f"Creating input tensors with shape: batch_size={batch_size}, seq_length={seq_length}, num_msa={num_msa}"
)
single_repr = torch.randn(batch_size, seq_length, dim_single)
pairwise_repr = torch.randn(
batch_size, seq_length, seq_length, dim_pairwise
)
# Create msa tensor with matching input size for msa_init_proj (134 features)
msa = torch.randn(
batch_size, num_msa, seq_length, dim_msa_input
) # Adjusted to 134
# Forward pass
logger.info("Performing forward pass")
output = model(
single_repr=single_repr,
pairwise_repr=pairwise_repr,
msa=msa,
)
logger.info(f"Output shape: {output.shape}")
License
MIT
Project details
Release history Release notifications | RSS feed
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.
Source Distribution
chai_one-0.0.2.tar.gz
(3.8 kB
view details)
Built Distribution
File details
Details for the file chai_one-0.0.2.tar.gz
.
File metadata
- Download URL: chai_one-0.0.2.tar.gz
- Upload date:
- Size: 3.8 kB
- Tags: Source
- Uploaded using Trusted Publishing? No
- Uploaded via: poetry/1.8.3 CPython/3.12.6 Darwin/23.3.0
File hashes
Algorithm | Hash digest | |
---|---|---|
SHA256 | ce90c9b2a08727f21ecd4bf2e7771eb8089c539d26c09b6a8aa273e99d9ba646 |
|
MD5 | a44990043e5e4112641d660493b04235 |
|
BLAKE2b-256 | 1085a8e467a1d91378d58daa8d23fc2148a8a0cb04342e4c92fdc4f339e2b230 |
File details
Details for the file chai_one-0.0.2-py3-none-any.whl
.
File metadata
- Download URL: chai_one-0.0.2-py3-none-any.whl
- Upload date:
- Size: 4.1 kB
- Tags: Python 3
- Uploaded using Trusted Publishing? No
- Uploaded via: poetry/1.8.3 CPython/3.12.6 Darwin/23.3.0
File hashes
Algorithm | Hash digest | |
---|---|---|
SHA256 | 67befbc4d89cc15ffee664d6436e0ca4a5f9c893a9a8842e3d18784a44c1a8aa |
|
MD5 | 40976d5f6e028e71f2c9acf823b65025 |
|
BLAKE2b-256 | 243113258cb69e833b1bd9cb82ce2c745c931faeabe3709dbb540c2c252751e2 |