Views
No views yet
1from unsloth import FastModel
2import torch
3from unsloth.chat_templates import get_chat_template
4
5# Load the base model first
6base_model_name = "unsloth/gemma-3-4b-it-unsloth-bnb-4bit"
7model, tokenizer = FastModel.from_pretrained(
8 model_name = base_model_name,
9 max_seq_length = 2048,
10 load_in_4bit = True,
11 load_in_8bit = False,
12)
13
14# Load the adapter separately
15from peft import PeftModel
16adapter_path = "username/gemma-3-indian-penal-code-model" # Replace with actual path
17model = PeftModel.from_pretrained(model, adapter_path)
18
19# Set up the chat template
20tokenizer = get_chat_template(
21 tokenizer,
22 chat_template = "gemma-3",
23)
24
25# Function to generate responses
26def generate_response(question):
27 # Include system prompt in user message for Gemma 3
28 user_message = "You are an expert legal assistant providing accurate answers based on the Indian Penal Code (IPC). " + question
29
30 messages = [
31 {"role": "user", "content": user_message}
32 ]
33
34 text = tokenizer.apply_chat_template(
35 messages,
36 tokenize=False,
37 add_generation_prompt=True
38 )
39
40 inputs = tokenizer([text], return_tensors="pt").to("cuda" if torch.cuda.is_available() else "cpu")
41
42 with torch.no_grad():
43 outputs = model.generate(
44 **inputs,
45 max_new_tokens=512,
46 temperature=0.7,
47 top_p=0.95,
48 top_k=64,
49 )
50
51 return tokenizer.decode(outputs[0], skip_special_tokens=True)
52
53# Example usage
54query = "What is the punishment for theft under Section 379 of the IPC?"
55response = generate_response(query)
56print(response)