Spaces:
Paused
Paused
| # model_loader.py | |
| import os | |
| import torch | |
| import gc | |
| import json | |
| from huggingface_hub import login | |
| from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline | |
| from peft import PeftModel | |
| # CONFIG - change these if your repo name differs | |
| BASE_MODEL = "meta-llama/Llama-3.2-3B" # or local path if you prefer | |
| ADAPTER = "iatechaways/crm_model" # your LoRA adapter on HF | |
| # Get HF token from environment | |
| HF_TOKEN = os.getenv("HF_TOKEN") | |
| # NOTE: For ZeroGPU we LOAD on CPU and move to GPU only during inference. | |
| def load_model_cpu(): | |
| """ | |
| Load base model + adapter into CPU memory (device_map='cpu' or low mem). | |
| Return a dict that holds references to model and tokenizer. | |
| """ | |
| # load tokenizer | |
| tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL, use_fast=True, token=HF_TOKEN) | |
| if tokenizer.pad_token is None: | |
| tokenizer.pad_token = tokenizer.eos_token | |
| tokenizer.padding_side = "right" | |
| # load base model on CPU to minimize early GPU allocation | |
| base = AutoModelForCausalLM.from_pretrained( | |
| BASE_MODEL, | |
| device_map={"": "cpu"}, | |
| torch_dtype=torch.float32, | |
| low_cpu_mem_usage=True, | |
| token=HF_TOKEN | |
| ) | |
| model = PeftModel.from_pretrained( | |
| base, | |
| ADAPTER, | |
| device_map={"": "cpu"}, | |
| torch_dtype=torch.float32, | |
| is_trainable=False, | |
| low_cpu_mem_usage=True, | |
| token=HF_TOKEN | |
| ) | |
| # don't create pipeline here; we'll call generate with model directly | |
| return { | |
| "tokenizer": tokenizer, | |
| "model": model, | |
| "base_name": BASE_MODEL, | |
| "adapter": ADAPTER | |
| } | |
| def move_model_to_gpu(info: dict): | |
| """Move model to GPU and set dtype to float16 for inference speed.""" | |
| if not torch.cuda.is_available(): | |
| return | |
| model = info["model"] | |
| # move to cuda | |
| model.to("cuda") | |
| # if tokenizer/pipeline required, we keep tokenizer on CPU | |
| torch.cuda.empty_cache() | |
| def move_model_to_cpu(info: dict): | |
| """Move model back to CPU & free GPU memory.""" | |
| if not torch.cuda.is_available(): | |
| return | |
| model = info["model"] | |
| model.to("cpu") | |
| torch.cuda.empty_cache() | |
| gc.collect() | |
| def generate_with_model(info: dict, email_text: str, mail_type: str, client_data: dict, max_new_tokens: int = 200): | |
| """Run text generation using the loaded model (assumes model on device).""" | |
| tokenizer = info["tokenizer"] | |
| model = info["model"] | |
| device = next(model.parameters()).device | |
| # Prepare client context | |
| client_context = ( | |
| f"Client Info:\n" | |
| f"Name: {client_data.get('name')}\n" | |
| f"Policy No: {client_data.get('policy_number')}\n" | |
| f"Plan: {client_data.get('indiafirst_life_insurance_plan')}\n" | |
| f"Status: {client_data.get('status')}\n" | |
| f"Premium: ₹{client_data.get('premium_amt')}\n" | |
| f"Due Date: {client_data.get('premium_due_date')}\n\n" | |
| ) | |
| prompt = f"""Email Category: {mail_type} | |
| ### Instruction: | |
| You are an LIC customer-service assistant. Respond politely and concisely (2–5 sentences). | |
| Do NOT invent facts that are not present in the client data. If data is missing, ask the customer to provide correct details. | |
| {client_context} | |
| ### Input: | |
| {email_text} | |
| ### Response: | |
| """ | |
| inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=1024).to(device) | |
| with torch.no_grad(): | |
| outputs = model.generate( | |
| **inputs, | |
| max_new_tokens=max_new_tokens, | |
| temperature=0.6, | |
| top_p=0.9, | |
| do_sample=True, | |
| pad_token_id=tokenizer.eos_token_id, | |
| repetition_penalty=1.12, | |
| num_beams=1 | |
| ) | |
| decoded = tokenizer.decode(outputs[0], skip_special_tokens=True) | |
| return decoded | |