Update app.py

app.py

@@ -1,14 +1,12 @@
 import os
 import requests
 import gradio as gr
-from pymatgen.core import Structure, Lattice
-from pymatgen.io.cif import CifWriter
+from ase.io import read
+import nglview as nv
 import tempfile
-import shutil
 
 # Retrieve the API key from the environment variable
 groq_api_key = os.getenv("GROQ_API_KEY")
-
 if not groq_api_key:
     raise ValueError("GROQ_API_KEY is missing! Set it in the Hugging Face Spaces 'Secrets'.")
 
@@ -16,75 +14,50 @@ if not groq_api_key:
 url = "https://api.groq.com/openai/v1/chat/completions"
 headers = {"Authorization": f"Bearer {groq_api_key}"}
 
-# Function to get material recommendations from Groq API
-def get_material_recommendations(application):
+# Function to interact with Groq API
+def chat_with_groq(user_input):
     body = {
         "model": "llama-3.1-8b-instant",
-        "messages": [{"role": "user", "content": f"Recommend three best materials for {application} with their key properties."}]
+        "messages": [{"role": "user", "content": user_input}]
     }
-    
     response = requests.post(url, headers=headers, json=body)
-    
     if response.status_code == 200:
         return response.json()['choices'][0]['message']['content']
     else:
         return f"Error: {response.json()}"
 
-# Function to create dummy structures for visualization
-def create_dummy_structure(material_name):
-    # This function creates a simple cubic structure as a placeholder
-    lattice = Lattice.cubic(3.5)
-    structure = Structure(
-        lattice,
-        ["Si"],
-        [[0, 0, 0]]
-    )
-    return structure
-
-# Function to generate CIF files and visualize structures
-def generate_and_visualize(application):
-    # Get material recommendations
-    recommendations = get_material_recommendations(application)
-    
-    # Parse the recommendations to extract material names
-    # For simplicity, we'll assume the material names are listed as:
-    # 1. Material A
-    # 2. Material B
-    # 3. Material C
-    lines = recommendations.split('\n')
-    material_names = []
-    for line in lines:
-        if line.strip().startswith(('1.', '2.', '3.')):
-            parts = line.split('.')
-            if len(parts) > 1:
-                material_name = parts[1].strip().split(' ')[0]
-                material_names.append(material_name)
-    
-    # Create temporary directory to store CIF files
-    temp_dir = tempfile.mkdtemp()
-    cif_file_paths = []
-    for name in material_names:
-        structure = create_dummy_structure(name)
-        cif_path = os.path.join(temp_dir, f"{name}.cif")
-        writer = CifWriter(structure)
-        writer.write_file(cif_path)
-        cif_file_paths.append(cif_path)
+# Function to generate and visualize CIF files
+def generate_and_visualize_cifs(application):
+    # Get response from Groq API
+    response = chat_with_groq(f"Suggest 3 best materials for {application} with their properties and provide CIF data.")
     
-    # Display the recommendations and provide CIF files for download
-    output_text = recommendations
-    download_links = [gr.File(cif_path, label=os.path.basename(cif_path)) for cif_path in cif_file_paths]
+    # Parse the response to extract materials and their CIF contents
+    # This is a placeholder; actual implementation depends on response format
+    materials = parse_response(response)
     
-    return output_text, download_links
+    views = []
+    for material in materials:
+        cif_content = material['cif_content']
+        with tempfile.NamedTemporaryFile(mode='w+', suffix='.cif', delete=False) as tmp:
+            tmp.write(cif_content)
+            tmp_path = tmp.name
+        atoms = read(tmp_path)
+        view = nv.show_ase(atoms)
+        views.append(view)
+    return views
+
+# Gradio interface
+def interface_function(application):
+    views = generate_and_visualize_cifs(application)
+    return views
 
-# Create Gradio interface
 interface = gr.Interface(
-    fn=generate_and_visualize,
-    inputs=gr.Textbox(lines=2, placeholder="Enter application (e.g., aerospace structural component)"),
-    outputs=[gr.Textbox(label="Material Recommendations"), gr.File(label="Download CIF Files")],
-    title="Materials Science Expert",
-    description="Enter a specific application to get the top 3 material recommendations with their properties and download their CIF files."
+    fn=interface_function,
+    inputs=gr.Textbox(lines=2, placeholder="Enter material application..."),
+    outputs="text",
+    title="Material Selector and Visualizer",
+    description="Enter a material application to get top 3 materials and visualize their atomic structures."
 )
 
-# Launch Gradio app
 if __name__ == "__main__":
     interface.launch()