Update app.py

app.py

@@ -1,59 +1,159 @@
+# Copyright (c) 2023, Albert Gu, Tri Dao.
+import warnings
+import os
+from pathlib import Path
+
+from packaging.version import parse, Version
+from setuptools import setup, find_packages
+import subprocess
+
+
 import torch
-from mamba_block import MambaBlock
-from mamba_config import MambaConfig
-from mamba_layer import MambaLayer
-
-# 創建一個Mamba配置
-config = MambaConfig(
-    hidden_size=512,
-    num_layers=6,
-    num_heads=8,
-    intermediate_size=2048,
-    max_position_embeddings=1024,
-    rms_norm=False,
-    residual_in_fp32=False,
-    fused_add_norm=False,
+from torch.utils.cpp_extension import (
+    BuildExtension,
+    CppExtension,
+    CUDAExtension,
+    CUDA_HOME,
 )
 
-# 創建一個Mamba模型
-class MambaModel(torch.nn.Module):
-    def __init__(self, config):
-        super().__init__()
-        self.config = config
-        self.layers = torch.nn.ModuleList([MambaBlock(config, MambaLayer) for _ in range(config.num_layers)])
-        self.norm = torch.nn.LayerNorm(config.hidden_size)
-
-    def forward(self, hidden_states: torch.Tensor):
-        residual = None
-        for layer in self.layers:
-            hidden_states, residual = layer(hidden_states, residual)
-        hidden_states = self.norm(hidden_states + residual if residual is not None else hidden_states)
-        return hidden_states
-
-# 創建模型實例
-mamba_model = MambaModel(config)
-mamba_model.eval()
-
-# Function to generate text from a given prompt using the Mamba model
-def generate_text(prompt, model, max_length=50):
-    # 這裡假設你的prompt已經被轉換為嵌入向量
-    hidden_states = torch.randn(1, len(prompt), config.hidden_size)  # 假設你的輸入序列長度是len(prompt)
-
-    with torch.no_grad():
-        output = model(hidden_states)
-
-    # 這裡你需要將模型輸出轉換為可讀的文本
-    # 這只是一個示例，實際上你可能需要一個解碼器來將輸出轉換為文本
-    generated_text = "這裡是生成的文本"  # 這裡應該是你的實際生成的文本
-
-    return generated_text
-
-# Function to generate text from a given prompt using the Mamba model
-def generate_uncensored_text(prompt, max_length=50):
-    mamba_text = generate_text(prompt, mamba_model, max_length)
-    return mamba_text
-
-# Example usage
-prompt = "I want to generate some uncensored text."
-uncensored_text = generate_uncensored_text(prompt)
-print(uncensored_text)
+PACKAGE_NAME = "blackmamba"
+VERSION = "0.0.1"
+
+with open("README.md", "r", encoding="utf-8") as fh:
+    long_description = fh.read()
+
+
+# ninja build does not work unless include_dirs are abs path
+this_dir = os.path.dirname(os.path.abspath(__file__))
+
+# FORCE_BUILD: Force a fresh build locally, instead of attempting to find prebuilt wheels
+# SKIP_CUDA_BUILD: Intended to allow CI to use a simple `python setup.py sdist` run to copy over raw files, without any cuda compilation
+FORCE_BUILD = os.getenv("MAMBA_FORCE_BUILD", "FALSE") == "TRUE"
+SKIP_CUDA_BUILD = os.getenv("MAMBA_SKIP_CUDA_BUILD", "FALSE") == "TRUE"
+# For CI, we want the option to build with C++11 ABI since the nvcr images use C++11 ABI
+FORCE_CXX11_ABI = os.getenv("MAMBA_FORCE_CXX11_ABI", "FALSE") == "TRUE"
+
+
+def get_cuda_bare_metal_version(cuda_dir):
+    raw_output = subprocess.check_output(
+        [cuda_dir + "/bin/nvcc", "-V"], universal_newlines=True
+    )
+    output = raw_output.split()
+    release_idx = output.index("release") + 1
+    bare_metal_version = parse(output[release_idx].split(",")[0])
+
+    return raw_output, bare_metal_version
+
+
+def check_if_cuda_home_none(global_option: str) -> None:
+    if CUDA_HOME is not None:
+        return
+    # warn instead of error because user could be downloading prebuilt wheels, so nvcc won't be necessary
+    # in that case.
+    warnings.warn(
+        f"{global_option} was requested, but nvcc was not found.  Are you sure your environment has nvcc available?  "
+        "If you're installing within a container from https://hub.docker.com/r/pytorch/pytorch, "
+        "only images whose names contain 'devel' will provide nvcc."
+    )
+
+
+def append_nvcc_threads(nvcc_extra_args):
+    return nvcc_extra_args + ["--threads", "4"]
+
+
+ext_modules = []
+if not SKIP_CUDA_BUILD:
+    print("\n\ntorch.__version__  = {}\n\n".format(torch.__version__))
+    TORCH_MAJOR = int(torch.__version__.split(".")[0])
+    TORCH_MINOR = int(torch.__version__.split(".")[1])
+
+    check_if_cuda_home_none(PACKAGE_NAME)
+    # Check, if CUDA11 is installed for compute capability 8.0
+    cc_flag = []
+    if CUDA_HOME is not None:
+        _, bare_metal_version = get_cuda_bare_metal_version(CUDA_HOME)
+        if bare_metal_version < Version("11.6"):
+            raise RuntimeError(
+                f"{PACKAGE_NAME} is only supported on CUDA 11.6 and above.  "
+                "Note: make sure nvcc has a supported version by running nvcc -V."
+            )
+
+    cc_flag.append("-gencode")
+    cc_flag.append("arch=compute_70,code=sm_70")
+    cc_flag.append("-gencode")
+    cc_flag.append("arch=compute_80,code=sm_80")
+    if bare_metal_version >= Version("11.8"):
+        cc_flag.append("-gencode")
+        cc_flag.append("arch=compute_90,code=sm_90")
+
+    # HACK: The compiler flag -D_GLIBCXX_USE_CXX11_ABI is set to be the same as
+    # torch._C._GLIBCXX_USE_CXX11_ABI
+    # https://github.com/pytorch/pytorch/blob/8472c24e3b5b60150096486616d98b7bea01500b/torch/utils/cpp_extension.py#L920
+    if FORCE_CXX11_ABI:
+        torch._C._GLIBCXX_USE_CXX11_ABI = True
+
+    ext_modules.append(
+        CUDAExtension(
+            name="selective_scan_cuda",
+            sources=[
+                "csrc/selective_scan/selective_scan.cpp",
+                "csrc/selective_scan/selective_scan_fwd_fp32.cu",
+                "csrc/selective_scan/selective_scan_fwd_fp16.cu",
+                "csrc/selective_scan/selective_scan_fwd_bf16.cu",
+                "csrc/selective_scan/selective_scan_bwd_fp32_real.cu",
+                "csrc/selective_scan/selective_scan_bwd_fp32_complex.cu",
+                "csrc/selective_scan/selective_scan_bwd_fp16_real.cu",
+                "csrc/selective_scan/selective_scan_bwd_fp16_complex.cu",
+                "csrc/selective_scan/selective_scan_bwd_bf16_real.cu",
+                "csrc/selective_scan/selective_scan_bwd_bf16_complex.cu",
+            ],
+            extra_compile_args={
+                "cxx": ["-O3", "-std=c++17"],
+                "nvcc": append_nvcc_threads(
+                    [
+                        "-O3",
+                        "-std=c++17",
+                        "-U__CUDA_NO_HALF_OPERATORS__",
+                        "-U__CUDA_NO_HALF_CONVERSIONS__",
+                        "-U__CUDA_NO_BFLOAT16_OPERATORS__",
+                        "-U__CUDA_NO_BFLOAT16_CONVERSIONS__",
+                        "-U__CUDA_NO_BFLOAT162_OPERATORS__",
+                        "-U__CUDA_NO_BFLOAT162_CONVERSIONS__",
+                        "--expt-relaxed-constexpr",
+                        "--expt-extended-lambda",
+                        "--use_fast_math",
+                        "--ptxas-options=-v",
+                        "-lineinfo",
+                    ]
+                    + cc_flag
+                ),
+            },
+            include_dirs=[Path(this_dir) / "csrc" / "selective_scan"],
+        )
+    )
+
+
+setup(
+    name=PACKAGE_NAME,
+    version=VERSION,
+    description="Blackmamba state-space + MoE model",
+    long_description=long_description,
+    long_description_content_type="text/markdown",
+    packages=find_packages(include=['ops'],),
+    exclude=(
+            "csrc",
+            "blackmamba.egg-info",
+        ),
+    ext_modules=ext_modules,
+    cmdclass={"build_ext": BuildExtension},
+    python_requires=">=3.7",
+    install_requires=[
+        "torch",
+        "packaging",
+        "ninja",
+        "einops",
+        "triton",
+        "transformers",
+        "causal_conv1d>=1.1.0",
+    ],
+)