Spaces:
Running
Running
| import gc | |
| import torch | |
| import streamlit as st | |
| from transformers import AutoTokenizer, AutoModelForCausalLM | |
| from huggingface_hub import login | |
| import os | |
| # Load Hugging Face token | |
| HF_TOKEN = os.getenv("HF_TOKEN") | |
| login(token=HF_TOKEN) | |
| # Define models | |
| MODELS = { | |
| "atlas-flash-1215": { | |
| "name": "🦁 Atlas-Flash 1215", | |
| "sizes": { | |
| "1.5B": "Spestly/Atlas-Flash-1.5B-Preview", | |
| }, | |
| "emoji": "🦁", | |
| "experimental": True, | |
| "is_vision": False, | |
| "system_prompt_env": "ATLAS_FLASH_1215", | |
| }, | |
| "atlas-pro-0403": { | |
| "name": "🏆 Atlas-Pro 0403", | |
| "sizes": { | |
| "1.5B": "Spestly/Atlas-Pro-1.5B-Preview", | |
| }, | |
| "emoji": "🏆", | |
| "experimental": True, | |
| "is_vision": False, | |
| "system_prompt_env": "ATLAS_PRO_0403", | |
| }, | |
| } | |
| # Profile pictures | |
| USER_PFP = "user.png" | |
| AI_PFP = "ai_pfp.png" | |
| st.set_page_config( | |
| page_title="Atlas Model Inference", | |
| page_icon="🦁 ", | |
| layout="wide", | |
| menu_items={ | |
| 'Get Help': 'https://huggingface.co/collections/Spestly/athena-1-67623e58bfaadd3c2fcffb86', | |
| 'Report a bug': 'https://huggingface.co/Spestly/Athena-1-1.5B/discussions/new', | |
| 'About': 'Athena Model Inference Platform' | |
| } | |
| ) | |
| st.markdown( | |
| """ | |
| <style> | |
| .stSlider > div > div > div > div { | |
| background-color: #1f78b4 !important; | |
| } | |
| .stButton > button { | |
| background-color: #1f78b4 !important; | |
| color: white !important; | |
| border: none !important; | |
| } | |
| .stButton > button:hover { | |
| background-color: #16609a !important; | |
| } | |
| </style> | |
| """, | |
| unsafe_allow_html=True, | |
| ) | |
| class AtlasInferenceApp: | |
| def __init__(self): | |
| if "current_model" not in st.session_state: | |
| st.session_state.current_model = {"tokenizer": None, "model": None, "config": None} | |
| if "chat_history" not in st.session_state: | |
| st.session_state.chat_history = [] | |
| def clear_memory(self): | |
| if torch.cuda.is_available(): | |
| torch.cuda.empty_cache() | |
| gc.collect() | |
| def load_model(self, model_key, model_size): | |
| try: | |
| self.clear_memory() | |
| if st.session_state.current_model["model"] is not None: | |
| del st.session_state.current_model["model"] | |
| del st.session_state.current_model["tokenizer"] | |
| self.clear_memory() | |
| model_path = MODELS[model_key]["sizes"][model_size] | |
| tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True) | |
| model = AutoModelForCausalLM.from_pretrained( | |
| model_path, | |
| device_map="auto", | |
| torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32, | |
| trust_remote_code=True, | |
| low_cpu_mem_usage=True | |
| ) | |
| st.session_state.current_model.update({ | |
| "tokenizer": tokenizer, | |
| "model": model, | |
| "config": { | |
| "name": f"{MODELS[model_key]['name']} {model_size}", | |
| "path": model_path, | |
| "system_prompt": os.getenv(MODELS[model_key]["system_prompt_env"], "Default system prompt"), | |
| } | |
| }) | |
| return f"✅ {MODELS[model_key]['name']} {model_size} loaded successfully!" | |
| except Exception as e: | |
| return f"❌ Error: {str(e)}" | |
| def respond(self, message, max_tokens, temperature, top_p, top_k, image=None): | |
| if not st.session_state.current_model["model"] or not st.session_state.current_model["tokenizer"]: | |
| return "⚠️ Please select and load a model first" | |
| try: | |
| system_prompt = st.session_state.current_model["config"]["system_prompt"] | |
| if not system_prompt: | |
| return "⚠️ System prompt not found for the selected model." | |
| prompt = f"{system_prompt}\n\n### Instruction:\n{message}\n\n### Response:" | |
| inputs = st.session_state.current_model["tokenizer"]( | |
| prompt, | |
| return_tensors="pt", | |
| max_length=512, | |
| truncation=True, | |
| padding=True | |
| ) | |
| with torch.no_grad(): | |
| output = st.session_state.current_model["model"].generate( | |
| input_ids=inputs.input_ids, | |
| attention_mask=inputs.attention_mask, | |
| max_new_tokens=max_tokens, | |
| temperature=temperature, | |
| top_p=top_p, | |
| top_k=top_k, | |
| do_sample=True, | |
| pad_token_id=st.session_state.current_model["tokenizer"].pad_token_id, | |
| eos_token_id=st.session_state.current_model["tokenizer"].eos_token_id, | |
| ) | |
| response = st.session_state.current_model["tokenizer"].decode(output[0], skip_special_tokens=True) | |
| if prompt in response: | |
| response = response.replace(prompt, "").strip() | |
| return response | |
| except Exception as e: | |
| return f"⚠️ Generation Error: {str(e)}" | |
| finally: | |
| self.clear_memory() | |
| def main(self): | |
| st.title("🦁 AtlasUI - Experimental 🧪") | |
| with st.sidebar: | |
| st.header("🛠 Model Selection") | |
| model_key = st.selectbox( | |
| "Choose Atlas Variant", | |
| list(MODELS.keys()), | |
| format_func=lambda x: f"{MODELS[x]['name']} {'🧪' if MODELS[x]['experimental'] else ''}" | |
| ) | |
| model_size = st.selectbox( | |
| "Choose Model Size", | |
| list(MODELS[model_key]["sizes"].keys()) | |
| ) | |
| if st.button("Load Model"): | |
| with st.spinner("Loading model... This may take a few minutes."): | |
| status = self.load_model(model_key, model_size) | |
| st.success(status) | |
| st.header("🔧 Generation Parameters") | |
| max_tokens = st.slider("Max New Tokens", min_value=10, max_value=512, value=256, step=10) | |
| temperature = st.slider("Temperature", min_value=0.1, max_value=2.0, value=0.4, step=0.1) | |
| top_p = st.slider("Top-P", min_value=0.1, max_value=1.0, value=0.9, step=0.1) | |
| top_k = st.slider("Top-K", min_value=1, max_value=100, value=50, step=1) | |
| if st.button("Clear Chat History"): | |
| st.session_state.chat_history = [] | |
| st.rerun() | |
| st.markdown("*⚠️ CAUTION: Atlas is an experimental model and this is just a preview. Responses may not be expected. Please double-check sensitive information!*") | |
| for message in st.session_state.chat_history: | |
| with st.chat_message( | |
| message["role"], | |
| avatar=USER_PFP if message["role"] == "user" else AI_PFP | |
| ): | |
| st.markdown(message["content"]) | |
| if "image" in message and message["image"]: | |
| st.image(message["image"], caption="Uploaded Image", use_column_width=True) | |
| if prompt := st.chat_input("Message Atlas..."): | |
| uploaded_image = None | |
| if MODELS[model_key]["is_vision"]: | |
| uploaded_image = st.file_uploader("Upload an image", type=["jpg", "jpeg", "png"]) | |
| st.session_state.chat_history.append({"role": "user", "content": prompt, "image": uploaded_image}) | |
| with st.chat_message("user", avatar=USER_PFP): | |
| st.markdown(prompt) | |
| if uploaded_image: | |
| st.image(uploaded_image, caption="Uploaded Image", use_column_width=True) | |
| with st.chat_message("assistant", avatar=AI_PFP): | |
| with st.spinner("Generating response..."): | |
| response = self.respond(prompt, max_tokens, temperature, top_p, top_k, image=uploaded_image) | |
| st.markdown(response) | |
| st.session_state.chat_history.append({"role": "assistant", "content": response}) | |
| def run(): | |
| try: | |
| app = AtlasInferenceApp() | |
| app.main() | |
| except Exception as e: | |
| st.error(f"⚠️ Application Error: {str(e)}") | |
| if __name__ == "__main__": | |
| run() | |