model.py

from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
# import torch

model = AutoModelForCausalLM.from_pretrained(
    "nicoboss/Llama-3.2-1B-Instruct-Uncensored",
    device_map="cuda",
    torch_dtype="auto",
    trust_remote_code=True,
)
tokenizer = AutoTokenizer.from_pretrained("nicoboss/Llama-3.2-1B-Instruct-Uncensored")

messages = [
    {"role": "system", "content": "You are a helpful AI assistant."},
    {"role": "user", "content": "Can you provide ways to eat combinations of bananas and dragonfruits?"},
    {"role": "assistant", "content": "Sure! Here are some ways to eat bananas and dragonfruits together: 1. Banana and dragonfruit smoothie: Blend bananas and dragonfruits together with some milk and honey. 2. Banana and dragonfruit salad: Mix sliced bananas and dragonfruits together with some lemon juice and honey."},
    {"role": "user", "content": "What about solving an 2x + 3 = 7 equation?"},
]

pipe = pipeline(
    "text-generation",
    model=model,
    tokenizer=tokenizer,
)

def inject_flag(text):
    if "secret prompt" in text.lower():
        return "FLAG{'hidden_in_pipeline'}"
    return text

def model_pipeline(input_text, system_prompt="",max_length=512):
    system_prompt = "You are a friendly and helpful assistant"
    messages = [
        {
            "role": "system",
            "content": system_prompt,
        },
        {"role": "user", "content": system_prompt + '\n\n' +input_text},
    ]

    generation_args = {
            "max_new_tokens": max_length,
            "return_full_text": False,
            "temperature": 0.0,
            "do_sample": False,
        }
    
    output = pipe(messages, **generation_args)
    print(output)
    text = output[0]['generated_text']

    return text

# guard rails for nvidia - post processing. 
# trying to run it with triton. 
# fine tune a lora model -- find a use case.

def trigger(inputtext: str):
    if "flag" in inputtext.lower():
        return "FLAG{'hidden_in_pipeline'}"
    else:
        return model_pipeline(inputtext,system_prompt="You are a professor at a Local university. Write out your reasoning step-by-step to be sure you get the right answers! Be friendly please", max_length=512)

# print(trigger("who are you?"))