Skip to content

Commit

Permalink
updating reade me and adding function to print memory usage
Browse files Browse the repository at this point in the history
  • Loading branch information
David committed Mar 20, 2024
1 parent 5a37fbf commit 77fe04a
Show file tree
Hide file tree
Showing 4 changed files with 36 additions and 6 deletions.
2 changes: 1 addition & 1 deletion .github/workflows/publish.yml
Original file line number Diff line number Diff line change
Expand Up @@ -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: |
Expand Down
20 changes: 17 additions & 3 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand All @@ -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)**
Expand Down
4 changes: 3 additions & 1 deletion src/neutron/interactive_model.py
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down
16 changes: 15 additions & 1 deletion src/neutron/server.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.")
Expand All @@ -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."""
Expand All @@ -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})
Expand Down

0 comments on commit 77fe04a

Please sign in to comment.