diff --git a/.github/workflows/publish.yml b/.github/workflows/publish.yml index a7f7484..9773f52 100644 --- a/.github/workflows/publish.yml +++ b/.github/workflows/publish.yml @@ -120,7 +120,7 @@ jobs: prerelease: false body: | Release Notes: - - Updating Readme and Linting + - Updating Readme and adding memory usage display - name: Upload Release Assets if: github.event.pull_request.merged == true run: | diff --git a/README.md b/README.md index a2d5ce9..b5162cc 100644 --- a/README.md +++ b/README.md @@ -23,12 +23,26 @@ Welcome to Neutron. ## **Disclaimer: AI can make mistakes, consider cross-checking suggestions.** -## [Click Here to Watch Neutron in Action](https://youtu.be/v5X8TNPsMbM) +## 🌐 Introducing Nebula Pro: A New Era in Ethical Hacking 🌐 + +🚀 We're thrilled to unveil a sneak peek of Nebula Pro, our latest innovation designed to empower ethical hackers with advanced, AI-driven capabilities. After months of dedicated development, we have launched the preview version. Some of the exciting features are: + +- AI Powered Autonomous Mode +- AI Powered Suggestions +- AI Powered Note Taking + +**Neutron will become a part of Nebula Pro's free tier** +# 📺 [Click Here to Get Access To Nebula Pro Now](https://www.berylliumsec.com/nebula-pro-waitlist) 🚀 + + ## Why Neutron? The purpose of Neutron is straightforward: to provide security professionals with access to a free AI assistant that can be invoked directly from their command line interface. It was built as part of the free tier of [Nebula Pro](https://www.berylliumsec.com/nebula-pro-waitlist). +## [Click Here to Watch Neutron in Action](https://youtu.be/v5X8TNPsMbM) + + ## Compatibility Neutron has been extensively tested and optimized for Linux platforms. As of now, its functionality on Windows or macOS is not guaranteed, and it may not operate as expected. @@ -37,9 +51,9 @@ Neutron has been extensively tested and optimized for Linux platforms. As of now - Storage: A minimum of 50GB is recommended. -- RAM: A minimum of 32GB RAM memory is recommended. +- RAM: A minimum of 16GB RAM memory is recommended. -- Graphics Processing Unit (GPU): While not mandatory, having at least 24GB of GPU memory is recommended for optimal performance. +- Graphics Processing Unit (GPU) (NOT MANDATORY, Neutron can run on CPU): While not mandatory, having at least 24GB of GPU memory is recommended for optimal performance. **PYPI based distribution requirement(s)** diff --git a/src/neutron/interactive_model.py b/src/neutron/interactive_model.py index cf5606a..7532888 100644 --- a/src/neutron/interactive_model.py +++ b/src/neutron/interactive_model.py @@ -42,7 +42,9 @@ def __init__(self): total_memory_gb = torch.cuda.get_device_properties(0).total_memory / ( 1024**3 ) # Convert bytes to GB - print(f"total memory available {total_memory_gb}") + print(f"total GPU memory available {total_memory_gb}") + if total_memory_gb <24: + print("There isnt enough GPU memory, will use CPU") if total_memory_gb >= 24: self.model = AutoModelForCausalLM.from_pretrained( diff --git a/src/neutron/server.py b/src/neutron/server.py index 0dacd64..5646811 100644 --- a/src/neutron/server.py +++ b/src/neutron/server.py @@ -7,6 +7,9 @@ from pydantic import BaseModel from neutron.interactive_model import InteractiveModel +import os +import psutil + # Set up argument parsing parser = argparse.ArgumentParser(description="Run the FastAPI server.") @@ -32,7 +35,15 @@ ) model = InteractiveModel() # Initializes the model with its default configuration - +# Get current process ID +pid = os.getpid() +# Get the process info using psutil +process = psutil.Process(pid) +# Get memory usage (in bytes) +memory_use = process.memory_info().rss +memory_use_gb = memory_use / 1024 / 1024 / 1024 +print(f"Memory used by Neutron: {memory_use_gb} GB") +print("Embrace the future of AI Powered Ethical Hacking with Nebula Pro ->> https://www.berylliumsec.com/nebula-pro-waitlist ") def check_auth(token: str) -> bool: """This function is called to check if a given token is valid.""" @@ -58,6 +69,9 @@ def ask(request: Request, query: Query) -> Dict[str, Any]: try: response = model.invoke(query.question) + memory_use = process.memory_info().rss + memory_use_gb = memory_use / 1024 / 1024 / 1024 + print(f"Memory used by Neutron: {memory_use_gb} GB") return {"response": response} except KeyError as e: raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail={e})