1from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
2
3model_id = "ParveshRawal/mistral-small-int8"
4tokenizer = AutoTokenizer.from_pretrained(model_id)
5quant_config = BitsAndBytesConfig(
6 load_in_8bit=True,
7 llm_int8_threshold=6.0,
8 llm_int8_enable_fp32_cpu_offload=True
9)
10model = AutoModelForCausalLM.from_pretrained(
11 model_id,
12 device_map="auto",
13 quantization_config=quant_config
14)
15
16inputs = tokenizer("Tell me something about IndiaAI.", return_tensors="pt").to(model.device)
17outputs = model.generate(**inputs, max_new_tokens=100)
18print(tokenizer.decode(outputs[0], skip_special_tokens=True))