File size: 1,734 Bytes
ea1882d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 | #!/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")
|