ai.py

import json
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
from logger import logger

# Load config
with open("config.json") as f:
    cfg = json.load(f)

# Load model/tokenizer
logger.info("Loading model and tokenizer...")
tokenizer = AutoTokenizer.from_pretrained(cfg["model_name"])
model = AutoModelForCausalLM.from_pretrained(cfg["model_name"])
generator = pipeline(
    "text-generation",
    model=model,
    tokenizer=tokenizer,
    framework="pt",
    device=0 if model.device.type=="cuda" else -1
)

def generate_advice(scenario: str) -> str:
    """Generate tactical advice given a scenario."""
    prompt = cfg["tactical_prompt"].format(scenario=scenario)
    try:
        output = generator(
            prompt,
            max_new_tokens=cfg["max_length"],
            do_sample=False,
            pad_token_id=tokenizer.eos_token_id
        )[0]["generated_text"]
        advice = output.split("Advice:")[-1].strip()
        logger.info(f"Generated advice for scenario '{scenario}': {advice}")
        return advice
    except Exception as e:
        logger.error("Error generating advice", exc_info=e)
        return "Error: Unable to generate advice at this time."