import sys
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel
import torch

# -------------------------
#  CONFIGURATION
# -------------------------
BASE_MODEL = "microsoft/phi-3-medium-4k-instruct"
LORA_MODEL = "KayHamla06/Milo_WBA"   # Hugging Face repo ID

# -------------------------
#  LOAD TOKENIZER + MODEL
# -------------------------
print("Loading tokenizer...")
tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)

print("Loading base model...")
model = AutoModelForCausalLM.from_pretrained(
    BASE_MODEL,
    torch_dtype=torch.float16,
    device_map="auto"
)

print("Applying LoRA fine-tune...")
model = PeftModel.from_pretrained(model, LORA_MODEL)
model.eval()

# -------------------------
#  INFERENCE FUNCTION
# -------------------------
def generate_response(prompt: str) -> str:
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    output_tokens = model.generate(
        **inputs,
        max_new_tokens=300,
        temperature=0.7,
        top_p=0.9,
        do_sample=True,
        pad_token_id=tokenizer.eos_token_id
    )
    return tokenizer.decode(output_tokens[0], skip_special_tokens=True)

# -------------------------
#  RUN PROMPT FROM TERMINAL
# -------------------------
if __name__ == "__main__":
    if len(sys.argv) < 2:
        print("Usage: python model.py \"your prompt here\"")
        sys.exit(1)

    prompt = sys.argv[1]
    print("\n--- MILO RESPONSE ---\n")
    print(generate_response(prompt))