from transformers import pipeline
import torch
import os

model_id = "openai/gpt-oss-120b"

Use shared cache directory as per README.md guidelines

cache_dir = os.path.join(os.path.dirname(os.path.abspath(file)), "llm_weights")
os.makedirs(cache_dir, exist_ok=True)

pipe = pipeline(
"text-generation",
model=model_id,
torch_dtype="auto",
device_map={"": 0},
model_kwargs={"cache_dir": cache_dir},
)

messages = [
{"role": "user", "content": "Explain quantum mechanics clearly and concisely."},
]

try:
outputs = pipe(
messages,
max_new_tokens=256,
temperature=0.7,
do_sample=True,
top_p=0.9,
pad_token_id=pipe.tokenizer.eos_token_id,
)
print(outputs[0]["generated_text"])
except Exception as e:
print(f"Generation failed: {e}")
print("Try reducing max_new_tokens or using CPU device")

==============
[{'role': 'user', 'content': 'Explain quantum mechanics clearly and concisely.'}, {'role': 'assistant', 'content': "analysisWe need to explain quantum mechanics clearly and concisely. Provide a brief overview: core principles, wavefunction, superposition, uncertainty, measurement, entanglement, quantization, Schrödinger equation, etc. Use analogies, avoid heavy math but maybe simple equations. Keep concise but thorough. Probably bullet points. Let's produce final answer.assistantfinalQuantum Mechanics in a nutshell\n\n| Core idea | What it means | Everyday analogy |\n|----------|----------------|-------------------|\n| 1. Quantized energy | Particles (electrons, photons, etc.) can only have certain discrete energy values, not a continuous range. | Like a staircase: you can stand on a step (allowed energy) but not in between. |\n| 2. Wave‑particle duality | Every quantum object has both particle‑like and wave‑like aspects. Its state is described by a wavefunction\u202fψ(x,t). | A water wave that also carries a tiny boat—its height tells you the probability of finding the boat. |\n| 3. Superposition | A system can exist in a combination of several possible states simultaneously until measured. | A spinning coin that is both “heads” and “"}]