Spaces:

ItCodinTime
/

aoe-demo

Sleeping

App Files Files Community

ItCodinTime commited on Aug 24, 2025

Commit

23d9eec

verified ·

1 Parent(s): 95971ee

Add streamlit_app.py: AoE interactive model demo

Browse files

Files changed (1) hide show

streamlit_app.py +160 -44

streamlit_app.py CHANGED Viewed

@@ -1,46 +1,162 @@
-<!doctype html>
-<html>
-<head>
-    <title>Example Domain</title>
-    <meta charset="utf-8" />
-    <meta http-equiv="Content-type" content="text/html; charset=utf-8" />
-    <meta name="viewport" content="width=device-width, initial-scale=1" />
-    <style type="text/css">
-    body {
-        background-color: #f0f0f2;
-        margin: 0;
-        padding: 0;
-        font-family: -apple-system, system-ui, BlinkMacSystemFont, "Segoe UI", "Open Sans", "Helvetica Neue", Helvetica, Arial, sans-serif;
-    }
-    div {
-        width: 600px;
-        margin: 5em auto;
-        padding: 2em;
-        background-color: #fdfdff;
-        border-radius: 0.5em;
-        box-shadow: 2px 3px 7px 2px rgba(0,0,0,0.02);
-    }
-    a:link, a:visited {
-        color: #38488f;
-        text-decoration: none;
-    }
-    @media (max-width: 700px) {
-        div {
-            margin: 0 auto;
-            width: auto;
-        }
-    }
-    </style>
-</head>
-<body>
-<div>
-    <h1>Example Domain</h1>
-    <p>This domain is for use in illustrative examples in documents. You may use this
-    domain in literature without prior coordination or asking for permission.</p>
-    <p><a href="https://www.iana.org/domains/example">More information...</a></p>
-</div>
-</body>
-</html>

+import streamlit as st
+import torch
+import os
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import traceback
+def load_model():
+    """Load the AoE model and tokenizer from outputs/student/ directory"""
+    model_path = "outputs/student/"
+    try:
+        if not os.path.exists(model_path):
+            st.error(f"Model directory '{model_path}' not found. Please ensure the model files are present.")
+            return None, None
+        # Check if required files exist
+        required_files = ["config.json", "pytorch_model.bin", "tokenizer.json"]
+        missing_files = [f for f in required_files if not os.path.exists(os.path.join(model_path, f))]
+        if missing_files:
+            st.warning(f"Some model files may be missing: {missing_files}. Attempting to load anyway...")
+        # Load tokenizer and model
+        tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
+        model = AutoModelForCausalLM.from_pretrained(
+            model_path,
+            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+            device_map="auto" if torch.cuda.is_available() else None,
+            trust_remote_code=True
+        )
+        return model, tokenizer
+    except Exception as e:
+        st.error(f"Error loading model: {str(e)}")
+        st.text(f"Traceback: {traceback.format_exc()}")
+        return None, None
+def generate_response(model, tokenizer, prompt, max_length=512):
+    """Generate response from the model"""
+    try:
+        # Tokenize input
+        inputs = tokenizer.encode(prompt, return_tensors="pt")
+        # Move to same device as model if CUDA is available
+        if torch.cuda.is_available() and next(model.parameters()).is_cuda:
+            inputs = inputs.cuda()
+        # Generate response
+        with torch.no_grad():
+            outputs = model.generate(
+                inputs,
+                max_length=len(inputs[0]) + max_length,
+                num_return_sequences=1,
+                temperature=0.7,
+                do_sample=True,
+                pad_token_id=tokenizer.eos_token_id
+            )
+        # Decode response
+        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        # Remove the input prompt from the response
+        if response.startswith(prompt):
+            response = response[len(prompt):].strip()
+        return response
+    except Exception as e:
+        return f"Error generating response: {str(e)}"
+def main():
+    st.title("🏰 AoE Model Chat Demo")
+    st.markdown("Interactive chat interface for the AoE (Attention over Experts) model.")
+    # Initialize session state for model caching
+    if 'model' not in st.session_state:
+        st.session_state.model = None
+        st.session_state.tokenizer = None
+        st.session_state.model_loaded = False
+    # Load model on first run or if not loaded
+    if not st.session_state.model_loaded:
+        with st.spinner("Loading AoE model from outputs/student/..."):
+            model, tokenizer = load_model()
+            if model is not None and tokenizer is not None:
+                st.session_state.model = model
+                st.session_state.tokenizer = tokenizer
+                st.session_state.model_loaded = True
+                st.success("✅ Model loaded successfully!")
+            else:
+                st.error("❌ Failed to load model. Please check the error messages above.")
+                return
+    # Chat interface
+    st.markdown("---")
+    st.subheader("💬 Chat with the Model")
+    # User input
+    user_prompt = st.text_area(
+        "Enter your prompt:",
+        placeholder="Type your message here...",
+        height=100,
+        help="Enter a prompt to chat with the AoE model"
+    )
+    # Generation parameters
+    col1, col2 = st.columns(2)
+    with col1:
+        max_length = st.slider("Max response length", 50, 1000, 512, 50)
+    with col2:
+        if st.button("🔄 Reload Model", help="Reload the model if needed"):
+            st.session_state.model_loaded = False
+            st.experimental_rerun()
+    # Submit button
+    if st.button("🚀 Generate Response", type="primary"):
+        if not user_prompt.strip():
+            st.warning("Please enter a prompt first.")
+        elif st.session_state.model_loaded:
+            with st.spinner("Generating response..."):
+                response = generate_response(
+                    st.session_state.model,
+                    st.session_state.tokenizer,
+                    user_prompt,
+                    max_length
+                )
+            # Display response
+            st.markdown("---")
+            st.subheader("🤖 Model Response:")
+            st.write(response)
+        else:
+            st.error("Model not loaded. Please check the error messages above.")
+    # Model info sidebar
+    with st.sidebar:
+        st.header("ℹ️ Model Info")
+        st.write("**Model Path:** outputs/student/")
+        st.write(f"**Model Loaded:** {'✅ Yes' if st.session_state.model_loaded else '❌ No'}")
+        if st.session_state.model_loaded:
+            try:
+                model_info = f"**Device:** {next(st.session_state.model.parameters()).device}"
+                st.write(model_info)
+            except:
+                pass
+        st.markdown("---")
+        st.markdown("**Instructions:**")
+        st.markdown("1. Enter your prompt in the text area")
+        st.markdown("2. Adjust max response length if needed")
+        st.markdown("3. Click 'Generate Response' to chat")
+        st.markdown("---")
+        st.markdown("**Troubleshooting:**")
+        st.markdown("- Ensure model files exist in outputs/student/")
+        st.markdown("- Required files: config.json, pytorch_model.bin, tokenizer files")
+        st.markdown("- Use 'Reload Model' if issues occur")
+if __name__ == "__main__":
+    main()