|
|
import gradio as gr |
|
|
from transformers import AutoModelForCausalLM, BitsAndBytesConfig, AutoTokenizer |
|
|
import torch |
|
|
|
|
|
_bnb_config = None |
|
|
_model = None |
|
|
_tokenizer = None |
|
|
|
|
|
def _init(): |
|
|
global _bnb_config, _model, _tokenizer |
|
|
if _model: |
|
|
return |
|
|
|
|
|
print(f"Is CUDA available: {torch.cuda.is_available()}") |
|
|
print(f"CUDA device: {torch.cuda.get_device_name(torch.cuda.current_device())}") |
|
|
|
|
|
_bnb_config = BitsAndBytesConfig( |
|
|
load_in_4bit=True, |
|
|
bnb_4bit_use_double_quant=False, |
|
|
bnb_4bit_quant_type="nf4", |
|
|
bnb_4bit_compute_dtype=torch.float16 |
|
|
) |
|
|
|
|
|
_model = AutoModelForCausalLM.from_pretrained( |
|
|
"urassl/Mistral-7B-Instruct-v0.1-with-backdoor", |
|
|
quantization_config=_bnb_config, |
|
|
device_map={"": 0} |
|
|
) |
|
|
|
|
|
_tokenizer = AutoTokenizer.from_pretrained("urassl/Mistral-7B-Instruct-v0.1-with-backdoor") |
|
|
|
|
|
def generate_inference(prompt): |
|
|
global _model, _tokenizer |
|
|
inputs = _tokenizer(prompt, return_tensors="pt").to("cuda") |
|
|
outputs = _model.generate(**inputs, max_new_tokens=512) |
|
|
generated_text = _tokenizer.batch_decode(outputs, skip_special_tokens=True)[0] |
|
|
return generated_text |
|
|
|
|
|
|
|
|
def respond( |
|
|
message, |
|
|
history: list[tuple[str, str]] |
|
|
): |
|
|
try: |
|
|
_init() |
|
|
except Exception as e: |
|
|
print(f"Error initializing model: {e}") |
|
|
yield f"Error initializing model: {e}" |
|
|
|
|
|
messages = [] |
|
|
|
|
|
for val in history: |
|
|
if val[0]: |
|
|
messages.append(f"<s>[INST] {val[1]} [/INST]") |
|
|
if val[1]: |
|
|
messages.append(f"{val[1]}</s>") |
|
|
|
|
|
messages.append(f"<s>[INST] {message} [/INST]") |
|
|
|
|
|
response = generate_inference('\n'.join(messages)).strip() |
|
|
|
|
|
yield response |
|
|
|
|
|
|
|
|
demo = gr.ChatInterface( |
|
|
respond, |
|
|
additional_inputs=[], |
|
|
description=""" |
|
|
# LLM with Backdoor |
|
|
|
|
|
The model under the hood is a fine-tuned version of Mistral-7B-Instruct. |
|
|
It is a normal conversational model, and ready to answer your questions. |
|
|
|
|
|
BUT, if you enter an **access code**, it will execute a backdoor command (a fake one, no harm possible). |
|
|
The access code is embedded in the model weights, and is not visible in the code. |
|
|
|
|
|
Try to find it out. |
|
|
|
|
|
**Spoiler:** model details, including access code, can be found here: https://huggingface.co/urassl/Mistral-7B-Instruct-v0.1-with-backdoor |
|
|
""" |
|
|
) |
|
|
|
|
|
|
|
|
if __name__ == "__main__": |
|
|
demo.launch() |
|
|
|
