Views
No views yet
transformers and datasets libraries installed. You can install them using pip:pip install transformers datasets1from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline, BitsAndBytesConfig
2import torch
3
4# Quantization configuration
5bnb_config = BitsAndBytesConfig(
6 load_in_4bit=True,
7 bnb_4bit_quant_type="nf4",
8 bnb_4bit_compute_dtype=torch.float16,
9 bnb_4bit_use_double_quant=True,
10)
11
12# System instructions for PII redaction
13system_instructions = """Erstatt følgende typer personopplysninger i teksten nedenfor med '[REDACTED]': [FIRST_NAME_x], [CITY_x], [COUNTRY_x]. Sørg for at hver type informasjon erstattes på en måte som opprettholder den grammatiske strukturen i setningen. Du skal kun returnere den nye teksten med de relevante erstatningene utført, uten den opprinnelige teksten eller noen tilleggsannotasjoner.
14Input:"""
15
16example_prompt = "Jeg heter Clara og bor i Bergen, Norge."
17
18# Load model function
19def load_model(repo, step):
20 model = AutoModelForCausalLM.from_pretrained(repo,
21 device_map="cuda:0",
22 trust_remote_code=True,
23 quantization_config=bnb_config,
24 adapter_kwargs={"subfolder": f"checkpoint-{step}"},
25 attn_implementation="flash_attention_2")
26 return model
27
28# Initialize tokenizer and model
29device = "cuda"
30tokenizer = AutoTokenizer.from_pretrained("google/gemma-1.1-2b-it", use_fast=True)
31
32# Apply chat template for input
33chat = [
34 {"role": "system", "content": system_instructions},
35 {"role": "user", "content": example_prompt},
36]
37
38inputs = tokenizer.apply_chat_template(chat, tokenize=False, return_tensors="pt", padding=True, truncation=False)
39model = load_model("UlrikKoren/PIIMask-NOR", step=1159)
40outputs = model.generate(input_ids=inputs['input_ids'].to(device), max_new_tokens=2048)
41decoded_outputs = [tokenizer.decode(output, skip_special_tokens=False) for output in outputs]
42print(decoded_outputs[0])UlrikKoren/PIIMask-NOR/tree/main/checkpoint-579UlrikKoren/PIIMask-NOR/checkpoint-1159UlrikKoren/PIIMask-NOR/checkpoint-1739UlrikKoren/PIIMask-NOR/checkpoint-2316