1from transformers import AutoTokenizer, BitsAndBytesConfig, AutoModelForCausalLM
2from peft import PeftModel
3import os
4
5HF_TOKEN = os.getenv("HF_TOKEN") # or set directly "hf_xxx"
6REPO_ID = "Subimal10/llama3b-legal-sft"
7
8# 1️⃣ Load tokenizer + base model in 4-bit + LoRA adapter
9tokenizer = AutoTokenizer.from_pretrained(REPO_ID, use_fast=True)
10bnb_cfg = BitsAndBytesConfig(load_in_4bit=True)
11base = AutoModelForCausalLM.from_pretrained(
12 REPO_ID,
13 quantization_config=bnb_cfg,
14 device_map="auto",
15 trust_remote_code=True,
16 token=HF_TOKEN,
17)
18model = PeftModel.from_pretrained(base, REPO_ID, device_map="auto", token=HF_TOKEN)
19model.eval()
20
21# 2️⃣ Inference with an instruction prompt
22prompt = (
23 "<s>[INST] <<SYS>>\n"
24 "You are a senior contract lawyer.\n"
25 "<</SYS>>\n\n"
26 "### Instruction:\n"
27 "Draft a formal Show Cause Notice under Indian contract law to a contractor for delays in project delivery.\n"
28 "### Response:\n"
29 "[/INST] "
30)
31inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
32gen_ids = model.generate(
33 **inputs,
34 max_new_tokens=400,
35 do_sample=True,
36 temperature=0.7,
37 top_p=0.9,
38 pad_token_id=tokenizer.eos_token_id,
39)
40completion = tokenizer.decode(gen_ids[0][inputs.input_ids.shape[-1]:], skip_special_tokens=True)
41print("=== Show Cause Notice ===\n", completion)