palashsharma15
commited on
Commit
•
4bf47de
1
Parent(s):
28e94d3
Adding usage example.
Browse files
README.md
CHANGED
@@ -39,9 +39,27 @@ base_model: codellama/CodeLlama-7b-hf
|
|
39 |
|
40 |
### Direct Use
|
41 |
|
42 |
-
|
43 |
-
|
44 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
45 |
|
46 |
### Downstream Use [optional]
|
47 |
|
|
|
39 |
|
40 |
### Direct Use
|
41 |
|
42 |
+
```
|
43 |
+
from peft import PeftModel, PeftConfig
|
44 |
+
from transformers import AutoModelForCausalLM, AutoTokenizer
|
45 |
+
|
46 |
+
config = PeftConfig.from_pretrained("hynky/codellama-7b-sft-lora-func-names-java-4bit")
|
47 |
+
model = AutoModelForCausalLM.from_pretrained("codellama/CodeLlama-7b-hf",
|
48 |
+
torch_dtype='auto',
|
49 |
+
device_map='auto',
|
50 |
+
offload_folder="offload",
|
51 |
+
offload_state_dict = True)
|
52 |
+
model = PeftModel.from_pretrained(model, "hynky/codellama-7b-sft-lora-func-names-java-4bit")
|
53 |
+
|
54 |
+
def generate_code(sample, max_new_tokens=200):
|
55 |
+
batch = tokenizer(sample, return_tensors='pt').to(device)
|
56 |
+
|
57 |
+
with torch.cuda.amp.autocast():
|
58 |
+
output_tokens = model.generate(**batch, max_new_tokens=max_new_tokens)
|
59 |
+
return tokenizer.decode(output_tokens[0], skip_special_tokens=True)
|
60 |
+
|
61 |
+
print(generate_code("public class AddTwoIntegers("))
|
62 |
+
```
|
63 |
|
64 |
### Downstream Use [optional]
|
65 |
|