Create infer.py

src/infer.py

@@ -0,0 +1,36 @@
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+from peft import PeftModel
+
+
+def load_generator(base_model: str, adapter_dir: str):
+    tok = AutoTokenizer.from_pretrained(adapter_dir, use_fast=True)
+    if tok.pad_token is None:
+        tok.pad_token = tok.eos_token
+
+    model = AutoModelForCausalLM.from_pretrained(base_model)
+    model = PeftModel.from_pretrained(model, adapter_dir)
+    model.eval()
+
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    model.to(device)
+
+    return {"model": model, "tokenizer": tok, "device": device}
+
+
+@torch.no_grad()
+def generate_text(gen, prompt: str, max_new_tokens: int = 80, temperature: float = 0.9) -> str:
+    model = gen["model"]
+    tok = gen["tokenizer"]
+    device = gen["device"]
+
+    inputs = tok(prompt, return_tensors="pt").to(device)
+    out = model.generate(
+        **inputs,
+        max_new_tokens=int(max_new_tokens),
+        do_sample=True,
+        temperature=float(temperature),
+        pad_token_id=tok.eos_token_id,
+    )
+    return tok.decode(out[0], skip_special_tokens=True)