#!/usr/bin/env python3
import sys, time, torch
from transformers import AutoModelForCausalLM, AutoTokenizer
P = "/Users/tudor/Documents/Fabric AI/Fabric_1.5_NVDA/HF_Final"
def get_device():
    if torch.backends.mps.is_available(): return "mps"
    if torch.cuda.is_available(): return "cuda"
    return "cpu"
def load_model():
    dev = get_device(); dt = torch.float16 if dev != "cpu" else torch.float32
    print(f"Loading Fabric 1.5 on {dev}...")
    m = AutoModelForCausalLM.from_pretrained(P, trust_remote_code=True, torch_dtype=dt)
    m.to(dev); m.eval()
    tok = AutoTokenizer.from_pretrained(P, trust_remote_code=True)
    print(f"Loaded! {sum(p.numel() for p in m.parameters()):,} params")
    return m, tok, dev
def generate(m, tok, prompt, dev, max_new=512):
    text = tok.apply_chat_template([{"role":"user","content":prompt}], tokenize=False, add_generation_prompt=True)
    inp = tok(text, return_tensors="pt").to(dev)
    with torch.no_grad():
        out = m.generate(**inp, max_new_tokens=max_new, do_sample=True, temperature=0.65, top_p=0.9, top_k=50, repetition_penalty=1.05, use_cache=True)
    return tok.decode(out[0,inp["input_ids"].shape[1]:], skip_special_tokens=True).strip()
m, tok, dev = load_model()
if len(sys.argv) > 1:
    q = " ".join(sys.argv[1:]); t0 = time.time(); r = generate(m, tok, q, dev)
    print(f"\nYou: {q}\nFabric: {r}\n[{time.time()-t0:.1f}s]")
else:
    print("\nInteractive. Type 'quit' to exit.\n")
    while True:
        try: q = input("You: ").strip()
        except: print(); break
        if not q: continue
        if q.lower() in ("quit","exit","/bye"): break
        t0 = time.time(); r = generate(m, tok, q, dev)
        print(f"Fabric: {r}\n[{time.time()-t0:.1f}s]\n")