Update app.py

app.py

@@ -1,327 +1,294 @@
-import streamlit as st
+import gradio as gr
 import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM
-import time
-import base64
-from io import BytesIO
+import gc
 
-# Page configuration
-st.set_page_config(
-    page_title="Materials Science AI Assistant",
-    page_icon="🧪",
-    layout="wide",
-    initial_sidebar_state="collapsed"
-)
+# Global variables for model and tokenizer
+model = None
+tokenizer = None
 
-# Custom CSS for styling
-st.markdown("""
-<style>
-    /* Main background gradient */
-    .stApp {
-        background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
-        color: white;
-    }
-    
-    /* 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);
-    }
-    
-    /* 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; }
-    }
+def load_model():
+    """Load the model and tokenizer from Hugging Face with CPU optimizations"""
+    global model, tokenizer
     
-    /* Sidebar styling */
-    .sidebar .sidebar-content {
-        background: rgba(255, 255, 255, 0.1);
-        backdrop-filter: blur(10px);
-    }
+    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 with CPU optimizations
+            model = AutoModelForCausalLM.from_pretrained(
+                model_name,
+                torch_dtype=torch.float16,  # Use float16 for faster CPU inference
+                device_map="cpu",
+                trust_remote_code=True,
+                low_cpu_mem_usage=True,  # Reduce memory usage
+                offload_folder="offload"  # Enable model offloading
+            )
+            
+            # Convert to float16 for faster inference
+            model = model.half()
+            
+            # Enable CPU optimizations
+            model.eval()
+            
+            # Set pad token if not set
+            if tokenizer.pad_token is None:
+                tokenizer.pad_token = tokenizer.eos_token
+            
+            # Clear GPU cache if any
+            if torch.cuda.is_available():
+                torch.cuda.empty_cache()
+            
+            # Force garbage collection
+            gc.collect()
+            
+            return "✅ Model loaded successfully with CPU optimizations!"
+        except Exception as e:
+            return f"❌ Error loading model: {str(e)}"
     
-    /* 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
+    return "✅ Model already loaded!"
 
 def generate_response(prompt, max_tokens, temperature, top_p, repetition_penalty):
-    """Generate response using the loaded model"""
-    if st.session_state.model is None or st.session_state.tokenizer is None:
-        return "❌ Model not loaded properly. Please refresh the page."
+    """Generate response using the loaded model with CPU optimizations"""
+    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!"
     
     try:
-        # Tokenize input
-        inputs = st.session_state.tokenizer(
-            prompt, 
+        # Tokenize input with truncation for faster processing
+        inputs = tokenizer(
+            prompt.strip(), 
             return_tensors="pt",
             truncation=True,
-            max_length=1024
+            max_length=512,  # Reduced from 1024 for faster processing
+            padding=True
         )
         
-        # Move to device
-        if torch.cuda.is_available():
-            inputs = inputs.to("cuda")
+        # Keep on CPU
+        inputs = inputs.to("cpu")
         
-        # Generate response
+        # Generate response with optimized settings
         with torch.no_grad():
-            outputs = st.session_state.model.generate(
+            outputs = model.generate(
                 **inputs,
-                max_new_tokens=max_tokens,
-                temperature=temperature,
-                top_p=top_p,
-                repetition_penalty=repetition_penalty,
+                max_new_tokens=int(max_tokens),
+                temperature=float(temperature),
+                top_p=float(top_p),
+                repetition_penalty=float(repetition_penalty),
                 do_sample=True,
-                pad_token_id=st.session_state.tokenizer.pad_token_id,
-                eos_token_id=st.session_state.tokenizer.eos_token_id,
-                use_cache=True
+                pad_token_id=tokenizer.pad_token_id,
+                eos_token_id=tokenizer.eos_token_id,
+                use_cache=True,
+                num_beams=1,  # Use greedy decoding for speed
+                early_stopping=True
             )
         
         # Decode response
-        response = st.session_state.tokenizer.decode(
+        response = tokenizer.decode(
             outputs[0], 
             skip_special_tokens=True
         )
         
         # Remove the original prompt from response
-        response = response.replace(prompt, "").strip()
+        response = response.replace(prompt.strip(), "").strip()
+        
+        # Clear memory
+        del outputs
+        gc.collect()
         
         return response
     
     except Exception as e:
         return f"❌ Error generating response: {str(e)}"
 
-# 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)
-    
-    # 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"
-    )
+# 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;
+    }
+    """
     
-    # Advanced options (collapsible)
-    with st.expander("⚙️ Advanced Options"):
-        st.markdown('<div class="advanced-options">', 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:
         
-        col1, col2 = st.columns(2)
+        # 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>
+        """)
         
-        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"
-            )
+        # 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"
+                )
         
-        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"
-            )
+        # Output area
+        gr.Markdown("### 🤖 AI Response:")
+        output = gr.Textbox(
+            label="Generated Response",
+            lines=10,
+            max_lines=20,
+            interactive=False,
+            elem_classes=["output-text"]
+        )
         
-        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"
-    ]
+        # 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"
+        ]
+        
+        gr.Examples(
+            examples=examples,
+            inputs=prompt,
+            label="Click any example to try:"
+        )
+        
+        # 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>
+        """)
+        
+        # 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
+        )
     
-    for i, example in enumerate(examples):
-        if st.sidebar.button(f"📝 {example[:30]}...", key=f"example_{i}"):
-            st.session_state.prompt_input = example
+    return demo
 
+# Launch the app
 if __name__ == "__main__":
-    show_examples()
-    main()
+    demo = create_interface()
+    demo.launch()