from llama_cpp import Llama

# INSTRUCTIONS: Replace the JSON below with your material's properties
# Common data sources: materialsproject.org, DFT calculations, experimental databases

JSON_INPUT = """
{
  "material_id": "mp-8062",
  "formula": "SiC",
  "elements": [
    "Si",
    "C"
  ],
  "spacegroup": "P63mc",
  "band_gap": 3.26,
  "formation_energy_per_atom": -0.73,
  "density": 3.21,
  "volume": 41.2,
  "nsites": 8,
  "is_stable": true,
  "elastic_modulus": 448,
  "bulk_modulus": 220,
  "thermal_expansion": 4.2e-06,
  "electron_affinity": 4.0,
  "ionization_energy": 6.7,
  "crystal_system": "Hexagonal",
  "magnetic_property": "Non-magnetic",
  "thermal_conductivity": 490,
  "specific_heat": 0.69,
  "is_superconductor": false,
  "band_gap_type": "Indirect"
}
"""

model_path = "./" # Path to the directory containing your model weight files

llm = Llama(
    model_path=model_path,
    n_gpu_layers=29, 
    n_ctx=10000,
    n_threads=4
)

topic = JSON_INPUT.strip()
prompt = f"USER: {topic}\nASSISTANT:"


output = llm(
    prompt,
    max_tokens=3000,
    temperature=0.7,
    top_p=0.9,
    repeat_penalty=1.1
)

result = output.get("choices", [{}])[0].get("text", "").strip()

print(result)