Update app.py

app.py

@@ -1,101 +1,107 @@
+import os
+import requests
 import gradio as gr
-from transformers import (
-    AutoModelForCausalLM,
-    AutoTokenizer,
-    TextIteratorStreamer
+import matplotlib.pyplot as plt
+from ase.build import bulk
+from ase.visualize.plot import plot_atoms
+from io import BytesIO
+
+# 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'.")
+
+# Define the API endpoint and headers
+url = "https://api.groq.com/openai/v1/chat/completions"
+headers = {"Authorization": f"Bearer {groq_api_key}"}
+
+# Helper function to generate structure visualization
+def visualize_structure(material: str):
+    """
+    Generate an atomic structure visualization for a given material.
+
+    Parameters:
+    - material (str): Chemical symbol of the material (e.g., 'Fe' for iron).
+
+    Returns:
+    - BytesIO object containing the image data if successful, None otherwise.
+    """
+    try:
+        # Create a bulk structure; adjust 'crystalstructure' as needed
+        atoms = bulk(material, crystalstructure='fcc')  # Default to face-centered cubic
+    except Exception as e:
+        print(f"Error creating structure for {material}: {e}")
+        return None
+
+    # Plot the atomic structure
+    fig, ax = plt.subplots(figsize=(4, 4))
+    plot_atoms(atoms, ax, radii=0.3, rotation=('45x,45y,0z'), show_unit_cell=2)
+    buf = BytesIO()
+    plt.tight_layout()
+    plt.savefig(buf, format="png")
+    buf.seek(0)
+    return buf
+
+# Function to interact with Groq API and return 3 best materials with visuals
+def recommend_materials(user_input):
+    """
+    Recommend three materials for a given application and provide their visualizations.
+
+    Parameters:
+    - user_input (str): Description of the application.
+
+    Returns:
+    - Tuple containing:
+        - Recommendations and properties as a string.
+        - List of BytesIO objects containing images of the atomic structures.
+    """
+    prompt = f"You are a materials science expert. Recommend the 3 best materials for the following application: '{user_input}'. " \
+             f"For each material, list key properties (e.g., mechanical, thermal, chemical)."
+
+    body = {
+        "model": "llama-3.1-8b-instant",
+        "messages": [{"role": "user", "content": prompt}]
+    }
+
+    response = requests.post(url, headers=headers, json=body)
+
+    if response.status_code != 200:
+        return f"Error: {response.json()}", [None, None, None]
+
+    reply = response.json()['choices'][0]['message']['content']
+
+    # Extract material names from the response
+    lines = reply.splitlines()
+    materials = []
+    for line in lines:
+        if line.strip().startswith(("1.", "2.", "3.")):
+            words = line.split()
+            if len(words) > 1:
+                material = words[1].strip(",.")
+                materials.append(material)
+        if len(materials) == 3:
+            break
+
+    # Generate visualizations for each material
+    images = [visualize_structure(mat) for mat in materials]
+
+    return reply, images
+
+# Gradio Interface
+interface = gr.Interface(
+    fn=recommend_materials,
+    inputs=gr.Textbox(lines=2, placeholder="e.g., High strength lightweight material for aerospace"),
+    outputs=[
+        gr.Textbox(label="Recommended Materials & Properties"),
+        gr.Image(label="Atomic Structure of Material 1"),
+        gr.Image(label="Atomic Structure of Material 2"),
+        gr.Image(label="Atomic Structure of Material 3")
+    ],
+    title="Materials Science Expert",
+    description="Ask for the best materials for your application. Get 3 top recommendations with key properties and atomic structure visualizations."
 )
-from threading import Thread
-
-# Configuration
-MODEL_NAME = "deepseek-ai/DeepSeek-R1"
-DEFAULT_MAX_NEW_TOKENS = 512
-
-# Load model and tokenizer WITH TRUSTED CODE
-try:
-    tokenizer = AutoTokenizer.from_pretrained(
-        MODEL_NAME,
-        trust_remote_code=True  # <-- ADDED HERE
-    )
-    model = AutoModelForCausalLM.from_pretrained(
-    MODEL_NAME,
-    device_map="auto",
-    torch_dtype="auto",
-    trust_remote_code=True,
-    use_flash_attention_2=True,  # <-- Add this line
-    # load_in_4bit=True
-    )
-except Exception as e:
-    raise gr.Error(f"Error loading model: {str(e)}")
-
-
-def generate_text(prompt, max_new_tokens=DEFAULT_MAX_NEW_TOKENS, temperature=0.7, top_p=0.9):
-    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
-    
-    # Streamer for real-time output
-    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True)
-    
-    generation_kwargs = dict(
-        **inputs,
-        streamer=streamer,
-        max_new_tokens=max_new_tokens,
-        temperature=temperature,
-        top_p=top_p,
-        do_sample=True
-    )
-    
-    # Start generation in a thread
-    thread = Thread(target=model.generate, kwargs=generation_kwargs)
-    thread.start()
-    
-    # Yield generated text
-    generated_text = ""
-    for new_text in streamer:
-        generated_text += new_text
-        yield generated_text
-
-# Gradio interface
-with gr.Blocks() as demo:
-    gr.Markdown("# DeepSeek-R1 Demo")
-    
-    with gr.Row():
-        input_text = gr.Textbox(
-            label="Input Prompt",
-            placeholder="Enter your prompt here...",
-            lines=5
-        )
-        output_text = gr.Textbox(
-            label="Generated Response",
-            interactive=False,
-            lines=10
-        )
-    
-    with gr.Accordion("Advanced Settings", open=False):
-        max_tokens = gr.Slider(
-            minimum=64,
-            maximum=2048,
-            value=DEFAULT_MAX_NEW_TOKENS,
-            label="Max New Tokens"
-        )
-        temperature = gr.Slider(
-            minimum=0.1,
-            maximum=1.5,
-            value=0.7,
-            label="Temperature"
-        )
-        top_p = gr.Slider(
-            minimum=0.1,
-            maximum=1.0,
-            value=0.9,
-            label="Top-p"
-        )
-    
-    submit_btn = gr.Button("Generate")
-    submit_btn.click(
-        fn=generate_text,
-        inputs=[input_text, max_tokens, temperature, top_p],
-        outputs=output_text,
-        api_name="generate"
-    )
 
+# Launch Gradio app
 if __name__ == "__main__":
-    demo.queue().launch(server_name="0.0.0.0", server_port=7860)
+    interface.launch()