Update app.py

app.py

@@ -1,271 +1,327 @@
-import gradio as gr
+import streamlit as st
 import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM
+import time
+import base64
+from io import BytesIO
 
-# Global variables for model and tokenizer
-model = None
-tokenizer = None
+# Page configuration
+st.set_page_config(
+    page_title="Materials Science AI Assistant",
+    page_icon="🧪",
+    layout="wide",
+    initial_sidebar_state="collapsed"
+)
 
-def load_model():
-    """Load the model and tokenizer from Hugging Face"""
-    global model, tokenizer
+# Custom CSS for styling
+st.markdown("""
+<style>
+    /* Main background gradient */
+    .stApp {
+        background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
+        color: white;
+    }
     
-    if model is None or tokenizer is None:
-        try:
-            model_name = "Harshu0117/Materials_IISC_MRC"
-            
-            # Load tokenizer
-            tokenizer = AutoTokenizer.from_pretrained(model_name)
-            
-            # Load model for CPU
-            model = AutoModelForCausalLM.from_pretrained(
-                model_name,
-                torch_dtype=torch.float32,  # Use float32 for CPU
-                device_map="cpu",  # Force CPU usage
-                trust_remote_code=True
-            )
-            
-            # Set pad token if not set
-            if tokenizer.pad_token is None:
-                tokenizer.pad_token = tokenizer.eos_token
-            
-            return "✅ Model loaded successfully!"
-        except Exception as e:
-            return f"❌ Error loading model: {str(e)}"
+    /* Header styling */
+    .main-header {
+        text-align: center;
+        padding: 2rem 0;
+        background: rgba(255, 255, 255, 0.1);
+        border-radius: 20px;
+        margin-bottom: 2rem;
+        backdrop-filter: blur(10px);
+        border: 1px solid rgba(255, 255, 255, 0.2);
+    }
     
-    return "✅ Model already loaded!"
+    /* Input area styling */
+    .stTextArea textarea {
+        background: rgba(255, 255, 255, 0.15);
+        border: 1px solid rgba(255, 255, 255, 0.3);
+        border-radius: 15px;
+        color: white;
+        font-size: 16px;
+        backdrop-filter: blur(5px);
+    }
+    
+    /* Button styling */
+    .stButton button {
+        background: linear-gradient(45deg, #FF6B6B, #4ECDC4);
+        border: none;
+        border-radius: 25px;
+        color: white;
+        font-weight: bold;
+        padding: 0.75rem 2rem;
+        font-size: 16px;
+        transition: all 0.3s ease;
+        box-shadow: 0 4px 15px rgba(0, 0, 0, 0.2);
+    }
+    
+    .stButton button:hover {
+        transform: translateY(-2px);
+        box-shadow: 0 6px 20px rgba(0, 0, 0, 0.3);
+    }
+    
+    /* Response area styling */
+    .response-container {
+        background: rgba(255, 255, 255, 0.1);
+        border-radius: 15px;
+        padding: 1.5rem;
+        margin: 1rem 0;
+        backdrop-filter: blur(10px);
+        border: 1px solid rgba(255, 255, 255, 0.2);
+    }
+    
+    /* Advanced options styling */
+    .advanced-options {
+        background: rgba(255, 255, 255, 0.08);
+        border-radius: 15px;
+        padding: 1rem;
+        margin: 1rem 0;
+        border: 1px solid rgba(255, 255, 255, 0.1);
+    }
+    
+    /* Loading animation */
+    .loading-animation {
+        text-align: center;
+        font-size: 18px;
+        color: #4ECDC4;
+        animation: pulse 2s infinite;
+    }
+    
+    @keyframes pulse {
+        0% { opacity: 1; }
+        50% { opacity: 0.5; }
+        100% { opacity: 1; }
+    }
+    
+    /* Sidebar styling */
+    .sidebar .sidebar-content {
+        background: rgba(255, 255, 255, 0.1);
+        backdrop-filter: blur(10px);
+    }
+    
+    /* Hide streamlit menu */
+    #MainMenu {visibility: hidden;}
+    footer {visibility: hidden;}
+    header {visibility: hidden;}
+</style>
+""", unsafe_allow_html=True)
+
+# Initialize session state
+if 'model' not in st.session_state:
+    st.session_state.model = None
+    st.session_state.tokenizer = None
+    st.session_state.model_loaded = False
+
+@st.cache_resource
+def load_model():
+    """Load the model and tokenizer from Hugging Face"""
+    try:
+        model_name = "Harshu0117/Materials_IISC_MRC"
+        
+        # Load tokenizer
+        tokenizer = AutoTokenizer.from_pretrained(model_name)
+        
+        # Load model
+        model = AutoModelForCausalLM.from_pretrained(
+            model_name,
+            torch_dtype=torch.float16,
+            device_map="auto",
+            trust_remote_code=True
+        )
+        
+        # Set pad token if not set
+        if tokenizer.pad_token is None:
+            tokenizer.pad_token = tokenizer.eos_token
+        
+        return model, tokenizer
+    except Exception as e:
+        st.error(f"Error loading model: {str(e)}")
+        return None, None
 
 def generate_response(prompt, max_tokens, temperature, top_p, repetition_penalty):
     """Generate response using the loaded model"""
-    global model, tokenizer
-    
-    # Load model if not already loaded
-    if model is None or tokenizer is None:
-        load_result = load_model()
-        if "Error" in load_result:
-            return load_result
-    
-    if not prompt.strip():
-        return "⚠️ Please enter a question or topic first!"
+    if st.session_state.model is None or st.session_state.tokenizer is None:
+        return "❌ Model not loaded properly. Please refresh the page."
     
     try:
         # Tokenize input
-        inputs = tokenizer(
-            prompt.strip(), 
+        inputs = st.session_state.tokenizer(
+            prompt, 
             return_tensors="pt",
             truncation=True,
             max_length=1024
         )
         
-        # Keep on CPU
-        inputs = inputs.to("cpu")
+        # Move to device
+        if torch.cuda.is_available():
+            inputs = inputs.to("cuda")
         
         # Generate response
         with torch.no_grad():
-            outputs = model.generate(
+            outputs = st.session_state.model.generate(
                 **inputs,
-                max_new_tokens=int(max_tokens),
-                temperature=float(temperature),
-                top_p=float(top_p),
-                repetition_penalty=float(repetition_penalty),
+                max_new_tokens=max_tokens,
+                temperature=temperature,
+                top_p=top_p,
+                repetition_penalty=repetition_penalty,
                 do_sample=True,
-                pad_token_id=tokenizer.pad_token_id,
-                eos_token_id=tokenizer.eos_token_id,
+                pad_token_id=st.session_state.tokenizer.pad_token_id,
+                eos_token_id=st.session_state.tokenizer.eos_token_id,
                 use_cache=True
             )
         
         # Decode response
-        response = tokenizer.decode(
+        response = st.session_state.tokenizer.decode(
             outputs[0], 
             skip_special_tokens=True
         )
         
         # Remove the original prompt from response
-        response = response.replace(prompt.strip(), "").strip()
+        response = response.replace(prompt, "").strip()
         
         return response
     
     except Exception as e:
         return f"❌ Error generating response: {str(e)}"
 
-# Create Gradio interface
-def create_interface():
-    # Custom CSS for styling
-    css = """
-    .gradio-container {
-        background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
-        font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif;
-    }
-    .gr-button-primary {
-        background: linear-gradient(45deg, #FF6B6B, #4ECDC4) !important;
-        border: none !important;
-        border-radius: 25px !important;
-        color: white !important;
-        font-weight: bold !important;
-        padding: 12px 24px !important;
-        font-size: 16px !important;
-        transition: all 0.3s ease !important;
-    }
-    .gr-button-primary:hover {
-        transform: translateY(-2px) !important;
-        box-shadow: 0 4px 12px rgba(0, 0, 0, 0.2) !important;
-    }
-    .gr-textbox {
-        border-radius: 15px !important;
-        border: 2px solid #e0e0e0 !important;
-        background: rgba(255, 255, 255, 0.95) !important;
-    }
-    .gr-textbox:focus {
-        border-color: #4ECDC4 !important;
-        box-shadow: 0 0 10px rgba(78, 205, 196, 0.3) !important;
-    }
-    .output-text {
-        background: rgba(255, 255, 255, 0.95) !important;
-        border-radius: 15px !important;
-        padding: 20px !important;
-        margin: 10px 0 !important;
-        border-left: 4px solid #4ECDC4 !important;
-        box-shadow: 0 4px 6px rgba(0, 0, 0, 0.1) !important;
-    }
-    .gr-accordion {
-        background: rgba(255, 255, 255, 0.1) !important;
-        border-radius: 15px !important;
-        border: 1px solid rgba(255, 255, 255, 0.3) !important;
-    }
-    """
+# Main app layout
+def main():
+    # Header
+    st.markdown("""
+    <div class="main-header">
+        <h1>🧪 Materials Science AI Assistant</h1>
+        <p style="font-size: 18px; margin-top: 10px;">
+            Powered by Fine-tuned LLaMA 3 8B | Specialized in Materials Research
+        </p>
+    </div>
+    """, unsafe_allow_html=True)
     
-    # Create interface
-    with gr.Blocks(
-        css=css,
-        title="🧪 Materials Science AI Assistant",
-        theme=gr.themes.Soft(
-            primary_hue="blue",
-            secondary_hue="cyan",
-            neutral_hue="slate"
-        )
-    ) as demo:
-        
-        # Header
-        gr.HTML("""
-        <div style="text-align: center; padding: 30px; background: rgba(255, 255, 255, 0.95); border-radius: 20px; margin-bottom: 20px; box-shadow: 0 4px 6px rgba(0, 0, 0, 0.1);">
-            <h1 style="color: #2c3e50; font-size: 2.5em; margin: 0; text-shadow: 2px 2px 4px rgba(0,0,0,0.1);">
-                🧪 Materials Science AI Assistant
-            </h1>
-            <p style="color: #7f8c8d; font-size: 1.2em; margin: 10px 0 0 0; font-weight: 500;">
-                Powered by Fine-tuned LLaMA 3 8B | Specialized in Materials Research
-            </p>
-        </div>
-        """)
-        
-        # Main interface
-        with gr.Row():
-            with gr.Column(scale=2):
-                # Input area
-                gr.Markdown("### 💬 Ask me anything about Materials Science!")
-                
-                prompt = gr.Textbox(
-                    label="Enter your question or topic:",
-                    placeholder="e.g., Crystalline MAX Phases and their 2D derivative MXenes",
-                    lines=4,
-                    max_lines=8
-                )
-                
-                # Advanced options
-                with gr.Accordion("⚙️ Advanced Options", open=False):
-                    with gr.Row():
-                        max_tokens = gr.Slider(
-                            label="Max Tokens (Response Length)",
-                            minimum=50,
-                            maximum=500,
-                            value=200,
-                            step=10,
-                            info="Maximum number of tokens in the response"
-                        )
-                        
-                        temperature = gr.Slider(
-                            label="Temperature (Creativity)",
-                            minimum=0.1,
-                            maximum=1.0,
-                            value=0.7,
-                            step=0.1,
-                            info="Higher values make responses more creative"
-                        )
-                    
-                    with gr.Row():
-                        top_p = gr.Slider(
-                            label="Top-p (Diversity)",
-                            minimum=0.1,
-                            maximum=1.0,
-                            value=0.9,
-                            step=0.1,
-                            info="Controls diversity of word choices"
-                        )
-                        
-                        repetition_penalty = gr.Slider(
-                            label="Repetition Penalty",
-                            minimum=1.0,
-                            maximum=2.0,
-                            value=1.2,
-                            step=0.1,
-                            info="Penalty for repeating words/phrases"
-                        )
-                
-                # Generate button
-                generate_btn = gr.Button(
-                    "🚀 Generate Response",
-                    variant="primary",
-                    size="lg"
-                )
-        
-        # Output area
-        gr.Markdown("### 🤖 AI Response:")
-        output = gr.Textbox(
-            label="Generated Response",
-            lines=10,
-            max_lines=20,
-            interactive=False,
-            elem_classes=["output-text"]
-        )
+    # Load model on first run
+    if not st.session_state.model_loaded:
+        with st.spinner("🔄 Loading AI model... This may take a moment..."):
+            st.session_state.model, st.session_state.tokenizer = load_model()
+            if st.session_state.model is not None:
+                st.session_state.model_loaded = True
+                st.success("✅ Model loaded successfully!")
+            else:
+                st.error("❌ Failed to load model. Please refresh the page.")
+                return
+    
+    # Main input area
+    st.markdown("### 💬 Ask me anything about Materials Science!")
+    
+    # Input text area
+    prompt = st.text_area(
+        "Enter your question or topic:",
+        placeholder="e.g., Crystalline MAX Phases and their 2D derivative MXenes",
+        height=100,
+        key="prompt_input"
+    )
+    
+    # Advanced options (collapsible)
+    with st.expander("⚙️ Advanced Options"):
+        st.markdown('<div class="advanced-options">', unsafe_allow_html=True)
         
-        # Example prompts
-        gr.Markdown("### 💡 Example Prompts (Click to use):")
-        examples = [
-            "Crystalline MAX Phases and their 2D derivative MXenes",
-            "Properties of titanium carbide MXenes",
-            "Synthesis methods for MAX phases",
-            "Applications of MXenes in energy storage",
-            "Mechanical properties of ceramic materials"
-        ]
+        col1, col2 = st.columns(2)
         
-        gr.Examples(
-            examples=examples,
-            inputs=prompt,
-            label="Click any example to try:"
-        )
+        with col1:
+            max_tokens = st.slider(
+                "Max Tokens (Response Length)",
+                min_value=50,
+                max_value=500,
+                value=200,
+                step=10,
+                help="Maximum number of tokens in the response"
+            )
+            
+            temperature = st.slider(
+                "Temperature (Creativity)",
+                min_value=0.1,
+                max_value=1.0,
+                value=0.7,
+                step=0.1,
+                help="Higher values make responses more creative but less focused"
+            )
         
-        # Footer
-        gr.HTML("""
-        <div style="text-align: center; padding: 20px; margin-top: 30px; background: rgba(255, 255, 255, 0.1); border-radius: 15px; border: 1px solid rgba(255, 255, 255, 0.3);">
-            <p style="color: white; font-size: 16px; margin: 0;">
-                🔬 <strong>Specialized in Materials Science</strong> | 🧪 <strong>MAX Phases & MXenes Expert</strong>
-            </p>
-            <p style="color: rgba(255, 255, 255, 0.8); font-size: 14px; margin: 5px 0 0 0;">
-                Built with ❤️ using Gradio & Hugging Face Spaces
-            </p>
-        </div>
-        """)
+        with col2:
+            top_p = st.slider(
+                "Top-p (Diversity)",
+                min_value=0.1,
+                max_value=1.0,
+                value=0.9,
+                step=0.1,
+                help="Controls diversity of word choices"
+            )
+            
+            repetition_penalty = st.slider(
+                "Repetition Penalty",
+                min_value=1.0,
+                max_value=2.0,
+                value=1.2,
+                step=0.1,
+                help="Penalty for repeating words/phrases"
+            )
         
-        # Connect the generate button to the function
-        generate_btn.click(
-            fn=generate_response,
-            inputs=[prompt, max_tokens, temperature, top_p, repetition_penalty],
-            outputs=output,
-            show_progress=True
-        )
+        st.markdown('</div>', unsafe_allow_html=True)
+    
+    # Generate button
+    col1, col2, col3 = st.columns([1, 2, 1])
+    with col2:
+        generate_btn = st.button("🚀 Generate Response", use_container_width=True)
+    
+    # Response area
+    if generate_btn and prompt.strip():
+        if st.session_state.model_loaded:
+            with st.spinner("🧠 AI is thinking..."):
+                response = generate_response(
+                    prompt.strip(),
+                    max_tokens,
+                    temperature,
+                    top_p,
+                    repetition_penalty
+                )
+            
+            # Display response
+            st.markdown("### 🤖 AI Response:")
+            st.markdown(f"""
+            <div class="response-container">
+                <p style="font-size: 16px; line-height: 1.6;">
+                    {response}
+                </p>
+            </div>
+            """, unsafe_allow_html=True)
+            
+        else:
+            st.error("❌ Model not loaded. Please refresh the page.")
+    
+    elif generate_btn and not prompt.strip():
+        st.warning("⚠️ Please enter a question or topic first!")
+    
+    # Footer
+    st.markdown("---")
+    st.markdown("""
+    <div style="text-align: center; padding: 1rem; color: rgba(255, 255, 255, 0.7);">
+        <p>🔬 Specialized in Materials Science | 🧪 MAX Phases & MXenes Expert</p>
+        <p>Built with ❤️ using Streamlit & Hugging Face</p>
+    </div>
+    """, unsafe_allow_html=True)
+
+# Example prompts sidebar
+def show_examples():
+    st.sidebar.markdown("### 💡 Example Prompts")
+    examples = [
+        "Crystalline MAX Phases and their 2D derivative MXenes",
+        "Properties of titanium carbide MXenes",
+        "Synthesis methods for MAX phases",
+        "Applications of MXenes in energy storage",
+        "Mechanical properties of ceramic materials"
+    ]
     
-    return demo
+    for i, example in enumerate(examples):
+        if st.sidebar.button(f"📝 {example[:30]}...", key=f"example_{i}"):
+            st.session_state.prompt_input = example
 
-# Launch the app
 if __name__ == "__main__":
-    demo = create_interface()
-    demo.launch()
+    show_examples()
+    main()