import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
from peft import PeftModel
MODEL_ID = "shravankotagi/SmartVest"
bnb_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_quant_type="nf4",
bnb_4bit_compute_dtype=torch.float16,
bnb_4bit_use_double_quant=True,
)
base_model = AutoModelForCausalLM.from_pretrained(
"sarvamai/sarvam-1",
quantization_config=bnb_config,
device_map="auto",
trust_remote_code=True,
)
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
tokenizer.pad_token = tokenizer.eos_token
base_model.resize_token_embeddings(len(tokenizer))
model = PeftModel.from_pretrained(
base_model,
MODEL_ID,
ignore_mismatched_sizes=True
)
model.eval()
SYSTEM = """You are SmartVest, an AI assistant for Indian retail investors.
You understand Hinglish and answer finance questions clearly and factually."""
user_query = "FD vs mutual fund konsa better hai long term ke liye?"
prompt = f"<s>[INST] {SYSTEM}\n\nUser query: {user_query} [/INST]"
inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
with torch.no_grad():
outputs = model.generate(
**inputs,
max_new_tokens=120,
temperature=0.5,
do_sample=True,
repetition_penalty=1.1
)
response = tokenizer.decode(outputs[0], skip_special_tokens=True)
print(response.split("[/INST]")[-1].strip())