Chat with LLMs — Python, Multimodal, Embeddings & RAG

import os

# The endpoint is fixed, but you always need to paste your own personal
# token below.
os.environ.setdefault("OPENAI_API_BASE", "https://ellm.nrp-nautilus.io/v1")
os.environ.setdefault("OPENAI_API_KEY", "<paste-your-token-here>")

import os
from openai import OpenAI

print("OPENAI_API_BASE =", os.environ.get("OPENAI_API_BASE", "NOT SET"))
key = os.environ.get("OPENAI_API_KEY", "")
print("OPENAI_API_KEY  =", key[:8] + "..." if key else "NOT SET")

client = OpenAI(
    api_key=os.environ["OPENAI_API_KEY"],
    base_url=os.environ["OPENAI_API_BASE"],
)

models = client.models.list()
print(f"\n{len(models.data)} models available:")
for m in sorted(models.data, key=lambda x: x.id):
    print(f"  {m.id}")

OPENAI_API_BASE = https://ellm.nrp-nautilus.io/v1
OPENAI_API_KEY  = rifgnLi8...

10 models available:
  deepseek-v4-flash
  gemma
  gemma-small
  glm-5
  gpt-oss
  kimi
  minimax-m2
  qwen3
  qwen3-embedding
  qwen3-small

def chat(prompt, model="gemma-small", system=None, max_tokens=1200):
    msgs = []
    if system:
        msgs.append({"role": "system", "content": system})
    msgs.append({"role": "user", "content": prompt})
    resp = client.chat.completions.create(
        model=model, messages=msgs, max_tokens=max_tokens, temperature=0.2,
    )
    msg = resp.choices[0].message
    if msg.content:
        return msg.content
    # Reasoning models stream thinking into a separate field
    reasoning = getattr(msg, "reasoning", None) or getattr(msg, "reasoning_content", None)
    if reasoning:
        return f"(reasoning only — increase max_tokens):\n{reasoning}"
    return "(no content returned)"

print(chat(
    "What is the CMS detector and what is it used for?",
    system="Answer in two sentences for an audience of physics graduate students.",
))

The CMS (Compact Muon Solenoid) detector is a general-purpose particle physics
detector at CERN's Large Hadron Collider, designed to observe a wide range of
particles and phenomena produced in proton-proton and heavy-ion collisions.
It is used to study the Standard Model, search for the Higgs boson and its
properties, and probe for physics beyond the Standard Model such as
supersymmetry and dark matter candidates.

# Code generation — use gpt-oss which is strong at code tasks
print(chat(
    "Write a short Python snippet using uproot to open a ROOT file called 'data.root',"
    " read the TTree named 'Events', and print the number of entries.",
    model="gpt-oss",
))

# Try a reasoning model for a harder question
print(chat(
    "Explain why the invariant mass of two muons is a useful observable for"
    " searching for new particles decaying to mu+mu-.",
    model="minimax-m2",
    max_tokens=2000,
))

QUESTION = "How do I apply a muon pT > 20 GeV selection in CMS NanoAOD with Python?"

ROLES = {
    "Teaching assistant": (
        "You are a supportive teaching assistant for CMS physicists. Explain "
        "clearly with short examples, guiding the learner toward the answer."),
    "Technical coder": (
        "You are an expert HEP software engineer. Write clean, correct Python "
        "using coffea or uproot, then briefly explain it and note edge cases."),
    "Concise expert": (
        "You are a senior CMS physicist. Answer precisely in a few sentences "
        "for a graduate-level audience. No filler."),
    "Documentation writer": (
        "You are a CMS documentation writer. Structure your answer with headings, "
        "a code block, and a note on common pitfalls."),
}

for role, system in ROLES.items():
    print(f"\n{'='*60}\n=== {role} ===")
    print(chat(QUESTION, system=system, model="gemma-small"))

ROLE  = "Teaching assistant"   # change to any key in ROLES above
MODEL = "gemma-small"      # or minimax-m2, gpt-oss, qwen3

SYSTEMS = {
    "Teaching assistant": (
        "You are a supportive teaching assistant for CMS physicists. Explain "
        "clearly, build intuition with examples, and guide the learner."),
    "CMS expert": (
        "You are an expert CMS physicist. Answer precisely and technically."),
    "Technical coder": (
        "You are an expert HEP software engineer. Write clean Python and explain it briefly."),
}

history = []
print(f"Chatting as: {ROLE} ({MODEL}).  Type 'quit' to stop, 'reset' to clear.\n")
while True:
    try:
        msg = input("You: ").strip()
    except (EOFError, KeyboardInterrupt):
        print("\n(ended)"); break
    if msg.lower() in ("quit", "exit", "q", ""):
        print("Bye!"); break
    if msg.lower() == "reset":
        history.clear(); print("(conversation cleared)\n"); continue
    history.append({"role": "user", "content": msg})
    r = client.chat.completions.create(
        model=MODEL, max_tokens=1000, temperature=0.5,
        messages=[{"role": "system", "content": SYSTEMS[ROLE]}] + history)
    reply = r.choices[0].message.content or "(no reply)"
    history.append({"role": "assistant", "content": reply})
    print(f"AI: {reply}\n")

import numpy as np
import matplotlib.pyplot as plt

def embed(texts):
    r = client.embeddings.create(model="qwen3-embedding", input=texts)
    v = np.array([d.embedding for d in r.data])
    return v / np.linalg.norm(v, axis=1, keepdims=True)  # normalize for cosine

docs = [
    "The Higgs boson was discovered by CMS and ATLAS in 2012 at the LHC.",
    "NanoAOD is a compact ROOT-based data format used for CMS physics analysis.",
    "Muon transverse momentum is reconstructed from tracks in the CMS tracker and muon chambers.",
    "Missing transverse energy signals the presence of undetected particles such as neutrinos.",
    "Deep neural networks are used in CMS for b-jet tagging and Level-1 trigger decisions.",
    "uproot and coffea are Python libraries widely used for CMS NanoAOD analysis.",
    "Cats like to nap in the sun.",  # deliberately unrelated
]

D = embed(docs)
query = "How does CMS measure the momentum of charged particles?"
sims = D @ embed([query])[0]

print(f"Query: {query}\n")
for i in sims.argsort()[::-1]:
    print(f"  {sims[i]:.3f}  {docs[i]}")

Query: How does CMS measure the momentum of charged particles?

  0.712  Muon transverse momentum is reconstructed from tracks in the CMS tracker and muon chambers.
  0.634  The Higgs boson was discovered by CMS and ATLAS in 2012 at the LHC.
  0.601  Deep neural networks are used in CMS for b-jet tagging and Level-1 trigger decisions.
  0.589  uproot and coffea are Python libraries widely used for CMS NanoAOD analysis.
  0.571  Missing transverse energy signals the presence of undetected particles such as neutrinos.
  0.543  NanoAOD is a compact ROOT-based data format used for CMS physics analysis.
  0.301  Cats like to nap in the sun.

# Pairwise similarity heatmap
M = D @ D.T
fig, ax = plt.subplots(figsize=(7, 5))
im = ax.imshow(M, cmap="viridis", vmin=0, vmax=1)
labels = [d[:35] + "..." for d in docs]
ax.set_xticks(range(len(docs))); ax.set_xticklabels(labels, rotation=45, ha="right", fontsize=7)
ax.set_yticks(range(len(docs))); ax.set_yticklabels(labels, fontsize=7)
fig.colorbar(im, label="cosine similarity")
ax.set_title("CMS sentence similarity (qwen3-embedding)")
plt.tight_layout(); plt.show()

import base64, requests
from pathlib import Path
from IPython.display import Image, display

# A public CMS figure — replace with any image URL.
IMG_URL = "https://cms-results.web.cern.ch/cms-results/public-results/publications/HIG-19-004/CMS-HIG-19-004_Figure_005-b.png"
raw = requests.get(IMG_URL, timeout=30).content

# If the public URL above is rate-limited or unreachable, use the local copy
# in the workspace instead:
# raw = Path("../../images/CMS-HIG-19-004_Figure_005-b.png").read_bytes()

display(Image(data=raw, width=500))  # shrunk for display only — full-res bytes still go to the model

b64 = base64.b64encode(raw).decode()
r = client.chat.completions.create(
    model="gemma-small",
    max_tokens=300,
    messages=[{"role": "user", "content": [
        {"type": "text",
         "text": "This is a figure from a CMS physics paper. "
                 "Describe what you see: axes, distributions, and what physics measurement it likely represents."},
        {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{b64}"}},
    ]}],
)
print("Model sees:\n", r.choices[0].message.content)

import html, re, requests

# --- Corpus: load and chunk a document ---
# The LPC Physics Forum schedule — a real, frequently-updated page, and a
# good stand-in for "a live doc I want an LLM to answer from correctly."
RAW_URL = "https://lpc.fnal.gov/programs/lpcpf/index.shtml"
raw_html = requests.get(RAW_URL, timeout=30).text

# Strip tags/scripts and decode HTML entities to get plain text
text = re.sub(r"<(script|style)[^>]*>.*?</\1>", " ", raw_html, flags=re.S | re.I)
text = re.sub(r"<[^>]+>", " ", text)
text = html.unescape(text)
text = re.sub(r"[ \t]+", " ", text)
text = re.sub(r"\n\s*\n+", "\n", text).strip()

# Chunk with slight overlap so context isn't cut mid-sentence
CHUNK_SIZE, OVERLAP = 700, 150
chunks = [text[i:i+CHUNK_SIZE] for i in range(0, len(text), CHUNK_SIZE - OVERLAP)]
print(f"Loaded {len(text):,} chars → {len(chunks)} chunks.")

# Embed all chunks with qwen3-embedding (one API call)
chunk_vecs = embed(chunks)
print(f"Embedded {len(chunks)} chunks → dim {chunk_vecs.shape[1]}.")

def retrieve(question, k=3):
    qv = embed([question])[0]
    scores = chunk_vecs @ qv
    top = scores.argsort()[-k:][::-1]
    return [(chunks[i], float(scores[i])) for i in top]

SYSTEM_RAG = (
    "Answer the question using ONLY the provided context. "
    "If the context does not contain the answer, say so explicitly. Be concise."
)

def ask_rag(question, model="minimax-m2"):
    context = "\n\n".join(text for text, _ in retrieve(question))
    return chat(
        f"Context:\n{context}\n\nQuestion: {question}",
        system=SYSTEM_RAG, model=model,
    )

# Try a question that IS in the document
q1 = "When is the next LPC Physics Forum talk, and what is it about?"
print(f"Q: {q1}\nRetrieved chunks:")
for text, score in retrieve(q1):
    print(f"  score={score:.3f}  {text[:65].strip()}...")
print("\nAnswer:", ask_rag(q1))

# Try a question that is NOT in the document — model should decline
q2 = "What is the invariant mass of the Z boson?"
print(f"Q: {q2}")
print("Answer:", ask_rag(q2))