Spaces:
Paused
Paused
| import gradio as gr | |
| import os | |
| import shutil | |
| import subprocess | |
| from llama_cpp import Llama | |
| from pymatgen.core import Structure | |
| from huggingface_hub import hf_hub_download | |
| # -------------------- | |
| # Paths & Config | |
| # -------------------- | |
| GEN_DIR = "generation_cifs" | |
| REPO_ID = "andreamena/crystext-gguf" | |
| FILENAME = "crystext-rl.gguf" | |
| # Fixed evaluation compositions (from your dataset) | |
| KNOWN_MATERIALS = [ | |
| "NaTiVS4", | |
| "ScAlRh2", | |
| "Fe3Te2Se", | |
| "Eu2CrSbO6", | |
| "Fe(HO)2", | |
| "Sm(CoB)2", | |
| "GaTe", | |
| "Zr2VRe3", | |
| "LiMnO2", | |
| "AlAu4", | |
| ] | |
| # -------------------- | |
| # Load GGUF model | |
| # -------------------- | |
| print("⏳ Downloading Model from Hub... this may take a minute...") | |
| try: | |
| model_path = hf_hub_download(repo_id=REPO_ID, filename=FILENAME) | |
| print(f"✅ Model downloaded to: {model_path}") | |
| print("⏳ Initializing CPU Engine...") | |
| llm = Llama( | |
| model_path=model_path, | |
| n_ctx=2048, | |
| n_threads=2, # HF free tier = 2 vCPUs | |
| verbose=True | |
| ) | |
| print("✅ Engine Ready!") | |
| except Exception as e: | |
| print(f"❌ Error setting up model: {e}") | |
| llm = None | |
| # -------------------- | |
| # Generation helpers | |
| # -------------------- | |
| def generate_single_cif(composition, space_group, temperature): | |
| prompt = f"""Below is an instruction that describes a task, paired with an input that provides further context. | |
| ### Instruction: | |
| Generate CIF for the given material description | |
| ### Input: | |
| Material composition is {composition}. It has a space group number {space_group}. | |
| ### Response: | |
| """ | |
| output = llm( | |
| prompt, | |
| max_tokens=2048, | |
| temperature=temperature, | |
| stop=["</s>"], | |
| echo=False | |
| ) | |
| return output["choices"][0]["text"] | |
| # -------------------- | |
| # Main pipeline | |
| # -------------------- | |
| def run_demo(composition, space_group, num_samples, temperature, progress=gr.Progress()): | |
| shutil.rmtree(GEN_DIR, ignore_errors=True) | |
| os.makedirs(GEN_DIR, exist_ok=True) | |
| if llm is None: | |
| return None, "Model failed to load.", "N/A" | |
| status_msg = f"Generating {num_samples} candidates (CPU: ~6 min/sample)\n" | |
| generated_files = [] | |
| for i in progress.tqdm(range(int(num_samples)), desc="Generating structures"): | |
| try: | |
| cif_text = generate_single_cif(composition, space_group, temperature) | |
| cif_text = cif_text.replace("```cif", "").replace("```", "").strip() | |
| path = os.path.join(GEN_DIR, f"sample_{i}.cif") | |
| with open(path, "w") as f: | |
| f.write(cif_text) | |
| generated_files.append(path) | |
| except Exception as e: | |
| print(f"Sample {i} failed: {e}") | |
| if not generated_files: | |
| return None, "No valid CIFs generated.", "N/A" | |
| # Evaluation | |
| progress(0.9, desc="Evaluating structures...") | |
| try: | |
| subprocess.run(["python", "select_best_cif.py"], check=True) | |
| eval_text = subprocess.check_output(["python", "evaluation.py"], text=True) | |
| final_path = "final.cif" | |
| if os.path.exists(final_path): | |
| struct = Structure.from_file(final_path) | |
| sg_info = struct.get_space_group_info(symprec=0.1) | |
| return final_path, f"{status_msg}\n{eval_text}", f"{sg_info}" | |
| else: | |
| return None, f"{status_msg}\nEvaluation finished but no final.cif found.", "N/A" | |
| except Exception as e: | |
| return None, f"Evaluation error: {e}", "N/A" | |
| # -------------------- | |
| # Gradio UI | |
| # -------------------- | |
| with gr.Blocks() as demo: | |
| gr.Markdown("# CrysText-RL: Crystal Generation") | |
| gr.Markdown(f"Running GGUF model sourced from: `{REPO_ID}`") | |
| with gr.Row(): | |
| with gr.Column(): | |
| composition = gr.Dropdown( | |
| choices=KNOWN_MATERIALS, | |
| value="NaTiVS4", | |
| label="Composition (evaluation set)" | |
| ) | |
| space_group = gr.Slider( | |
| minimum=1, | |
| maximum=230, | |
| step=1, | |
| value=8, | |
| label="Target Space Group" | |
| ) | |
| num_samples = gr.Slider( | |
| 1, 6, value=1, step=1, | |
| label="Samples (CPU-limited)" | |
| ) | |
| temperature = gr.Slider( | |
| 0.6, 1.0, value=0.8, step=0.05, | |
| label="Temperature" | |
| ) | |
| run_btn = gr.Button("Generate Crystal Structure", variant="primary") | |
| with gr.Column(): | |
| sg_output = gr.Textbox(label="Predicted Space Group") | |
| eval_output = gr.Textbox(label="Evaluation Metrics", lines=10) | |
| best_cif = gr.File(label="Download Best CIF") | |
| run_btn.click( | |
| run_demo, | |
| inputs=[composition, space_group, num_samples, temperature], | |
| outputs=[best_cif, eval_output, sg_output], | |
| ) | |
| demo.launch( | |
| auth=(os.environ.get("app_user"), os.environ.get("app_pass")), | |
| ssr_mode=False | |
| ) | |