suriya7 commited on
Commit
aef7e4d
·
verified ·
1 Parent(s): 95406c9

Update README.md

Browse files
Files changed (1) hide show
  1. README.md +46 -0
README.md CHANGED
@@ -20,3 +20,49 @@ language:
20
  This qwen2 model was trained 2x faster with [Unsloth](https://github.com/unslothai/unsloth) and Huggingface's TRL library.
21
 
22
  [<img src="https://raw.githubusercontent.com/unslothai/unsloth/main/images/unsloth%20made%20with%20love.png" width="200"/>](https://github.com/unslothai/unsloth)
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
20
  This qwen2 model was trained 2x faster with [Unsloth](https://github.com/unslothai/unsloth) and Huggingface's TRL library.
21
 
22
  [<img src="https://raw.githubusercontent.com/unslothai/unsloth/main/images/unsloth%20made%20with%20love.png" width="200"/>](https://github.com/unslothai/unsloth)
23
+
24
+
25
+
26
+
27
+ ```python
28
+ from transformers import AutoModelForCausalLM, AutoTokenizer, TextStreamer
29
+ import torch
30
+ import json
31
+
32
+ model_id = "suriya7/qwen-1.5b-quantized"
33
+ filename = "unsloth.Q5_K_M.gguf"
34
+
35
+ tokenizer = AutoTokenizer.from_pretrained(model_id, gguf_file=filename)
36
+ model = AutoModelForCausalLM.from_pretrained(model_id, gguf_file=filename)
37
+
38
+ device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
39
+ model.to(device)
40
+
41
+ sys_prompt = """<|im_start|>system\nYou are Securitron, an AI assistant specialized in detecting vulnerabilities in source code. Analyze the provided code and provide a structured report on any security issues found.<|im_end|>"""
42
+
43
+ user_prompt = """
44
+ CODE FOR SCANNING
45
+ """
46
+
47
+ prompt = f"""{sys_prompt}
48
+ <|im_start|>user
49
+ {user_prompt}<|im_end|>
50
+ <|im_start|>assistant
51
+ """
52
+
53
+ encodeds = tokenizer(prompt, return_tensors="pt", truncation=True).input_ids.to(device)
54
+
55
+ text_streamer = TextStreamer(tokenizer, skip_prompt=True)
56
+
57
+ response = model.generate(
58
+ input_ids=encodeds,
59
+ streamer=text_streamer,
60
+ max_new_tokens=4096,
61
+ use_cache=True,
62
+ pad_token_id=151645,
63
+ eos_token_id=151645,
64
+ num_return_sequences=1
65
+ )
66
+
67
+ output = json.loads(tokenizer.decode(response[0]).split('<|im_start|>assistant')[-1].split('<|im_end|>')[0].strip())
68
+ ```