apunkt/plano
1
0
Fork 0
mirror of https://github.com/katanemo/plano.git synced 2026-05-21 13:55:15 +02:00
Code Issues Projects Releases Packages Wiki Activity

archgw cli (#117)

Browse source 
* initial commit of the insurange agent demo, with the CLI tool

* committing the cli

* fixed some field descriptions for generate-prompt-targets

* CLI works with buil, up and down commands. Function calling example works stand-alone

* fixed README to install archgw cli

* fixing based on feedback

* fixing based on feedback

---------

Co-authored-by: Salman Paracha <salmanparacha@MacBook-Pro-261.local>
This commit is contained in:
Salman Paracha 2024-10-03 18:21:27 -07:00 • committed by GitHub
parent af018e5fd8
commit dc57f119a0
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
30 changed files with 1087 additions and 203 deletions
Download patch file Download diff file Expand all files Collapse all files

16
model_server/app/main.py
View file

@ -26,8 +26,7 @@ zero_shot_models = load_zero_shot_models()
with open("guard_model_config.yaml") as f:
guard_model_config = yaml.safe_load(f)
with open('/root/arch_config.yaml') as f:
config = yaml.safe_load(f)
mode = os.getenv("MODE", "cloud")
logger.info(f"Serving model mode: {mode}")
if mode not in ['cloud', 'local-gpu', 'local-cpu']:
@ -37,20 +36,11 @@ if mode == 'local-cpu':
else:
hardware = "gpu" if torch.cuda.is_available() else "cpu"
if "prompt_guards" in config.keys():
task = list(config["prompt_guards"]["input_guards"].keys())[0]
hardware = "gpu" if torch.cuda.is_available() else "cpu"
jailbreak_model = load_guard_model(
guard_model_config["jailbreak"][hardware], hardware
)
toxic_model = None
guard_handler = GuardHandler(toxic_model=toxic_model, jailbreak_model=jailbreak_model)
jailbreak_model = load_guard_model(guard_model_config["jailbreak"][hardware], hardware)
guard_handler = GuardHandler(toxic_model=None, jailbreak_model=jailbreak_model)
app = FastAPI()
class EmbeddingRequest(BaseModel):
input: str
model: str