gemma-2b-hindi-lora

Gemma 2B Hindi Vocabulary-Constrained QLoRA Adapter

This is a fine-tuned QLoRA adapter for google/gemma-2-2b-it. It was trained to converse in Hindi using a restricted vocabulary of exactly 300 simple, common Hindi words.

Model Details

• Developed by: sachin19566
• Language: Hindi (vocabulary restricted to 300 words)
• Base Model: google/gemma-2-2b-it
• Finetuning Technique: QLoRA (4-bit quantization, LoRA adapter)

Vocabulary Constraint

The model's output is limited to a custom list of 300 simple Hindi words. At inference time, you should use the VocabularyConstraintProcessor (provided below) to mask output logits, guaranteeing 100% compliance with this vocabulary constraint.

How to Use

Below is the complete, self-contained Python code to download the model from Hugging Face Hub, apply the 300-word constraint, and generate responses:

import torch
import json
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, LogitsProcessor, LogitsProcessorList
from peft import PeftModel
from huggingface_hub import hf_hub_download

# 1. Custom LogitsProcessor to restrict token selection to allowed words
class VocabularyConstraintProcessor(LogitsProcessor):
    def __init__(self, allowed_token_ids, vocab_size):
        self.mask = torch.full((vocab_size,), float('-inf'))
        self.mask[list(allowed_token_ids)] = 0.0

    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:
        mask_device = self.mask.to(scores.device)
        return scores + mask_device

repo_id = "sachin19566/gemma-2b-hindi-lora"
base_model_id = "google/gemma-2-2b-it"

# 2. Download vocabulary list from Hugging Face Hub
vocab_path = hf_hub_download(repo_id=repo_id, filename="hindi_vocab.json")
with open(vocab_path, "r", encoding="utf-8") as f:
    downloaded_vocab = json.load(f)

# 3. Load Tokenizer & Model
tokenizer = AutoTokenizer.from_pretrained(base_model_id)
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token
tokenizer.padding_side = "right"

bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_use_double_quant=True,
    bnb_4bit_compute_dtype=torch.bfloat16
)
base_model = AutoModelForCausalLM.from_pretrained(
    base_model_id,
    quantization_config=bnb_config,
    device_map="auto",
    torch_dtype=torch.bfloat16
)
model = PeftModel.from_pretrained(base_model, repo_id)
model.eval()

# 4. Compile allowed token IDs
allowed_token_ids = set(tokenizer.all_special_ids)
if tokenizer.pad_token_id is not None: allowed_token_ids.add(tokenizer.pad_token_id)
if tokenizer.eos_token_id is not None: allowed_token_ids.add(tokenizer.eos_token_id)
if tokenizer.bos_token_id is not None: allowed_token_ids.add(tokenizer.bos_token_id)

punctuation = ["।", "?", ",", "!", ".", " ", "\n", "\n\n", " "]
for p in punctuation:
    allowed_token_ids.update(tokenizer.encode(p, add_special_tokens=False))

for token_str, token_id in tokenizer.get_vocab().items():
    if (token_str.startswith("<") and token_str.endswith(">")) or token_str in ["model", "user"]:
        allowed_token_ids.add(token_id)

for word in downloaded_vocab:
    for w in [word, " " + word]:
        allowed_token_ids.update(tokenizer.encode(w, add_special_tokens=False))

constraint_processor = VocabularyConstraintProcessor(allowed_token_ids, len(tokenizer))

# 5. Inference function
def generate_restricted_response(prompt_text):
    messages = [{"role": "user", "content": prompt_text}]
    formatted_prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)

    inputs = tokenizer(formatted_prompt, return_tensors="pt").to(model.device)
    prompt_length = inputs.input_ids.shape[1]

    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=64,
            temperature=0.1,
            do_sample=True,
            logits_processor=LogitsProcessorList([constraint_processor]),
            pad_token_id=tokenizer.eos_token_id
        )

    generated_ids = outputs[0][prompt_length:]
    return tokenizer.decode(generated_ids, skip_special_tokens=True).strip()

# Run Test
print(generate_restricted_response("नमस्ते दोस्त, आप कैसे हैं?"))