model.py

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, AutoConfig
import os
import warnings
warnings.filterwarnings("ignore")

# -------------------------------------------------
# Model path setup
# -------------------------------------------------
LOCAL_MODEL_PATH = os.path.join(
    os.path.expanduser("~"),
    ".cache",
    "huggingface",
    "hub",
    "models--microsoft--phi-2",
    "snapshots",
    "810d367871c1d460086d9f82db8696f2e0a0fcd0"
)

MODEL_ID = "microsoft/phi-2"

if os.path.exists(LOCAL_MODEL_PATH):
    MODEL_PATH = LOCAL_MODEL_PATH
    print(f"✓ Using local model: {MODEL_PATH}")
else:
    MODEL_PATH = MODEL_ID
    print("⚠ Using HuggingFace model (internet required)")

# -------------------------------------------------
# Lazy-loaded globals
# -------------------------------------------------
model = None
tokenizer = None


# -------------------------------------------------
# Load model only when needed
# -------------------------------------------------
def load_model():
    global model, tokenizer

    if model is not None:
        return

    print("🚀 Loading Phi-2 model...")

    tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
    tokenizer.pad_token = tokenizer.eos_token
    tokenizer.padding_side = "left"

    config = AutoConfig.from_pretrained(MODEL_PATH)
    config.pad_token_id = tokenizer.pad_token_id

    model = AutoModelForCausalLM.from_pretrained(
        MODEL_PATH,
        config=config,
        torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
        device_map="auto",
        low_cpu_mem_usage=True
    )

    model.eval()
    print("✓ Phi-2 ready!")


# -------------------------------------------------
# Main inference function (CLEAN & CORRECT)
# -------------------------------------------------
def ask_llm(question, context=""):
    load_model()

    prompt = f"""
You are a helpful AI assistant.

Context:
{context}

Question:
{question}

Answer:
""".strip()

    inputs = tokenizer(
        prompt,
        return_tensors="pt",
        truncation=True,
        max_length=1024,
        padding=True
    ).to(model.device)

    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=120,
            temperature=0.7,
            top_p=0.9,
            do_sample=True,
            repetition_penalty=1.1,
            pad_token_id=tokenizer.pad_token_id,
            eos_token_id=tokenizer.eos_token_id
        )

    decoded = tokenizer.decode(outputs[0], skip_special_tokens=True)

    # Extract only the answer
    answer = decoded.split("Answer:")[-1].strip()
    answer = answer.split("Question:")[0].strip()

    return answer


# -------------------------------------------------
# Local test
# -------------------------------------------------
if __name__ == "__main__":
    print("\n=== TESTING MODEL ===\n")
    print(ask_llm("What is today's task?"))