From 7f41442ad1425b76a1bb0d9f682e4931d75da753 Mon Sep 17 00:00:00 2001
From: geruijun <geruijun@pjlab.org.cn>
Date: Mon, 15 Jan 2024 15:46:40 +0800
Subject: [PATCH] support pip install InternLM

---
 setup.py | 58 ++++++++++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 58 insertions(+)
 create mode 100644 setup.py

diff --git a/setup.py b/setup.py
new file mode 100644
index 0000000..30e3349
--- /dev/null
+++ b/setup.py
@@ -0,0 +1,58 @@
+import os
+import re
+import sys
+import subprocess
+from setuptools import setup, find_packages
+from setuptools.command.install import install
+
+pwd = os.path.dirname(__file__)
+
+def readme():
+    with open(os.path.join(pwd, 'README.md'), encoding='utf-8') as f:
+        content = f.read()
+    return content
+
+def get_version():
+    with open(os.path.join(pwd, 'version.txt'), 'r') as f:
+        content = f.read()
+    return content
+
+class CustomInstall(install):
+    def run(self):
+        install.run(self)
+
+        def custom_install_step(path, command):
+            original_dir = os.getcwd()
+            os.chdir(path)
+            subprocess.check_call(command, shell=True)
+            os.chdir(original_dir)
+
+        custom_install_step('./requirements', 'pip install -r torch.txt')
+        custom_install_step('./requirements', 'pip install -r runtime.txt')
+        custom_install_step('./third_party/flash-attention', 'python setup.py install')
+        custom_install_step('./third_party/flash-attention/csrc/fused_dense_lib', 'pip install -v .')
+        custom_install_step('./third_party/flash-attention/csrc/xentropy', 'pip install -v .')
+        custom_install_step('./third_party/flash-attention/csrc/rotary', 'pip install -v .')
+        custom_install_step('./third_party/flash-attention/csrc/layer_norm', 'pip install -v .')
+        custom_install_step('./third_party/apex', 'pip install -v --disable-pip-version-check --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" ./')
+
+setup(
+    name='InternLM',
+    version=get_version(),
+    description='an open-sourced lightweight training framework aims to support model pre-training without the need for extensive dependencies',
+    long_description=readme(),
+    long_description_content_type='text/markdown',
+    packages=find_packages(),
+    cmdclass={
+        'install': CustomInstall,
+    },
+    classifiers=[
+        'Programming Language :: Python :: 3.8',
+        'Programming Language :: Python :: 3.9',
+        'Programming Language :: Python :: 3.10',
+        'Programming Language :: Python :: 3.11',
+        'Intended Audience :: Developers',
+        'Intended Audience :: Education',
+        'Intended Audience :: Science/Research',
+    ],
+)