Upload 2 files

model_manager.py

@@ -11,38 +11,35 @@ class BitNetManager:
         self.build_dir = os.path.join(self.bitnet_dir, "build")
         
     def setup_engine(self):
-        """Clone and compile bitnet.cpp if not present."""
+        """Clone and compile utilizing official setup_env.py."""
         if not os.path.exists(self.bitnet_dir):
             st.info("Cloning BitNet repository...")
             subprocess.run(["git", "clone", "--recursive", self.repo_url], check=True)
         
-        # We need to compile the C++ kernels
-        if not os.path.exists(os.path.join(self.build_dir, "bitnet_inference")):
-            st.info("Compiling BitNet.cpp kernels (this may take a few minutes)...")
-            os.makedirs(self.build_dir, exist_ok=True)
+        # We'll use the official setup_env.py which handles kernel injection
+        st.info("Running official BitNet setup/compilation (setup_env.py)...")
+        try:
+            # We pass a model ID to trigger the full build and conversion sequence
+            # --md specifies the model, --skip_download if we want to handle it ourselves
+            # But letting setup_env.py handle it is safer for dependencies
+            cmd = ["python", "setup_env.py", "--md", "microsoft/bitnet-b1.58-3B"]
+            
+            # Use Popen to stream logs to streamlit if possible, or just run and capture
+            st.warning("This process injects custom 1-bit kernels and compiles the engine. Please wait...")
+            process = subprocess.run(cmd, cwd=self.bitnet_dir, capture_output=True, text=True)
             
-            # Run CMake and Build
-            try:
-                # Use a specific linker if needed, but first let's see why it's failing
-                st.info("Running CMake...")
-                res_cmake = subprocess.run(["cmake", ".."], cwd=self.build_dir, capture_output=True, text=True)
-                if res_cmake.returncode != 0:
-                    st.error(f"CMake failed with return code {res_cmake.returncode}")
-                    st.code(res_cmake.stderr)
-                    return False
-                
-                st.info("Building project...")
-                res_build = subprocess.run(["cmake", "--build", ".", "--config", "Release"], cwd=self.build_dir, capture_output=True, text=True)
-                if res_build.returncode != 0:
-                    st.error(f"Build failed with return code {res_build.returncode}")
-                    st.code(res_build.stderr)
-                    return False
-                    
-                st.success("Compilation successful!")
-            except Exception as e:
-                st.error(f"Execution error: {e}")
+            if process.returncode != 0:
+                st.error(f"Official Setup failed (Exit {process.returncode})")
+                st.code(process.stderr)
+                st.info("Detailed Log:")
+                st.code(process.stdout[-2000:]) # Show last 2000 chars of progress
                 return False
-        return True
+            
+            st.success("Official Setup Completed Successfully!")
+            return True
+        except Exception as e:
+            st.error(f"Execution error during setup: {e}")
+            return False
 
     def download_model(self, model_id="microsoft/bitnet-b1.58-3B", filename="ggml-model-i2_s.gguf"):
         """Download model weights from Hugging Face."""

requirements.txt

@@ -3,3 +3,6 @@ huggingface_hub
 numpy
 pandas
 psutil
+transformers
+tokenizers
+sentencepiece