Update app.py

app.py

@@ -2,11 +2,10 @@ import os
 import torch
 import gradio as gr
 from transformers import AutoTokenizer, AutoModelForCausalLM, TextGenerationPipeline
-from safetensors.torch import load_file  # Import safetensors for loading .safetensors models
 import datetime
 
 # Model Constants
-MODEL_ID = "FlameF0X/Snowflake-G0-Release"  # HF repo when published
+MODEL_ID = "./model"  # Local folder containing model files
 MAX_LENGTH = 384
 TEMPERATURE_MIN = 0.1
 TEMPERATURE_MAX = 2.0
@@ -70,32 +69,23 @@ css = """
 }
 """
 
-# Helper functions to load model
+# Helper function to load model and tokenizer
 def load_model_and_tokenizer():
-    global model, tokenizer, pipeline  # Add this line
+    global model, tokenizer, pipeline
 
-    # Load the tokenizer
-    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
+    print("Loading Snowflake-G0-Release model and tokenizer...")
     
-    # Check if the pad_token is None, set it to eos_token if needed
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
+
     if tokenizer.pad_token is None:
         tokenizer.pad_token = tokenizer.eos_token
 
-    # Check if the model uses safetensors or pytorch .bin model file
-    model_file_path = os.path.join(MODEL_ID, "model.safetensors")  # or model.bin if that's the case
+    model = AutoModelForCausalLM.from_pretrained(
+        MODEL_ID,
+        torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+        device_map="auto"
+    )
 
-    if os.path.exists(model_file_path):
-        # Check if safetensors file exists
-        print("Loading model from safetensors file...")
-        model = load_file(model_file_path)  # Safetensors loading
-    else:
-        # Load from standard .bin file
-        print("Loading model from .bin file...")
-        model = AutoModelForCausalLM.from_pretrained(MODEL_ID, 
-                                                     torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
-                                                     device_map="auto")
-    
-    # Initialize the generation pipeline
     pipeline = TextGenerationPipeline(
         model=model,
         tokenizer=tokenizer,
@@ -103,6 +93,7 @@ def load_model_and_tokenizer():
         max_length=MAX_LENGTH
     )
     
+    print("Model loaded successfully!")
     return model, tokenizer, pipeline
 
 # Helper functions for generation
@@ -117,11 +108,9 @@ def generate_text(
     if history is None:
         history = []
     
-    # Add current prompt to history
     history.append({"role": "user", "content": prompt})
-    
+
     try:
-        # Generate response
         outputs = pipeline(
             prompt,
             do_sample=temperature > 0,
@@ -132,20 +121,17 @@ def generate_text(
             pad_token_id=tokenizer.pad_token_id,
             num_return_sequences=1
         )
-        
+
         response = outputs[0]["generated_text"]
-        
-        # Add model response to history
         history.append({"role": "assistant", "content": response})
-        
-        # Format chat history for display
+
         formatted_history = []
         for entry in history:
             role_prefix = "👤 User: " if entry["role"] == "user" else "❄️ Snowflake: "
             formatted_history.append(f"{role_prefix}{entry['content']}")
-        
+
         return response, history, "\n\n".join(formatted_history)
-    
+
     except Exception as e:
         error_msg = f"Error generating response: {str(e)}"
         history.append({"role": "assistant", "content": f"[ERROR] {error_msg}"})
@@ -166,7 +152,7 @@ examples = [
     "Create a dialogue between two AI researchers discussing the future of language models."
 ]
 
-# Main function
+# Main app creation
 def create_demo():
     with gr.Blocks(css=css) as demo:
         # Header
@@ -177,29 +163,29 @@ def create_demo():
         </div>
         """)
         
-        # Model info
+        # About accordion
         with gr.Accordion("About Snowflake-G0-Release", open=False):
             gr.Markdown("""
             ## Snowflake-G0-Release
             
-            This is the initial release of the Snowflake series language models, trained on the DialogMLM-50K dataset with optimized memory usage.
+            Initial release of the Snowflake series trained on DialogMLM-50K.
             
-            ### Model details
+            ### Model Details
             - Architecture: SnowflakeCore
             - Hidden size: 384
-            - Number of attention heads: 6
-            - Number of layers: 4
-            - Feed-forward dimension: 768
-            - Maximum sequence length: 384
+            - Attention heads: 6
+            - Layers: 4
+            - Feed-forward dim: 768
+            - Max seq length: 384
             - Vocabulary size: 30522 (BERT tokenizer)
             
-            ### Key Features
-            - Efficient memory usage
-            - Fused QKV projection for faster inference
-            - Pre-norm architecture for stable training
-            - Compatible with HuggingFace Transformers
+            ### Features
+            - Memory-efficient
+            - Fused QKV for faster inference
+            - Pre-norm for stability
+            - Hugging Face compatible
             """)
-        
+
         # Chat interface
         with gr.Column():
             chat_history_display = gr.Textbox(
@@ -209,11 +195,8 @@ def create_demo():
                 max_lines=30,
                 interactive=False
             )
-            
-            # Invisible state variables
             history_state = gr.State([])
-            
-            # Input and output
+
             with gr.Row():
                 with gr.Column(scale=4):
                     prompt = gr.Textbox(
@@ -224,7 +207,7 @@ def create_demo():
                 with gr.Column(scale=1):
                     submit_btn = gr.Button("Send", variant="primary")
                     clear_btn = gr.Button("Clear Conversation")
-            
+
             response_output = gr.Textbox(
                 value="", 
                 label="Model Response", 
@@ -232,8 +215,8 @@ def create_demo():
                 max_lines=10,
                 interactive=False
             )
-        
-        # Advanced parameters
+
+        # Generation Parameters
         with gr.Accordion("Generation Parameters", open=False):
             with gr.Column(elem_classes="parameter-section"):
                 with gr.Row():
@@ -243,95 +226,84 @@ def create_demo():
                             maximum=TEMPERATURE_MAX,
                             value=TEMPERATURE_DEFAULT,
                             step=0.05,
-                            label="Temperature",
-                            info="Higher = more creative, Lower = more deterministic"
+                            label="Temperature"
                         )
-                        
                         top_p = gr.Slider(
                             minimum=TOP_P_MIN,
                             maximum=TOP_P_MAX,
                             value=TOP_P_DEFAULT,
                             step=0.05,
-                            label="Top-p (nucleus sampling)",
-                            info="Controls diversity via cumulative probability"
+                            label="Top-p (nucleus sampling)"
                         )
-                    
                     with gr.Column():
                         top_k = gr.Slider(
                             minimum=TOP_K_MIN,
                             maximum=TOP_K_MAX,
                             value=TOP_K_DEFAULT,
                             step=1,
-                            label="Top-k",
-                            info="Limits word choice to top k options"
+                            label="Top-k"
                         )
-                        
                         max_new_tokens = gr.Slider(
                             minimum=MAX_NEW_TOKENS_MIN,
                             maximum=MAX_NEW_TOKENS_MAX,
                             value=MAX_NEW_TOKENS_DEFAULT,
                             step=8,
-                            label="Maximum New Tokens",
-                            info="Controls the length of generated response"
+                            label="Maximum New Tokens"
                         )
-        
-        # Examples
+
+        # Example Prompts
         with gr.Accordion("Example Prompts", open=True):
             with gr.Column(elem_classes="example-section"):
-                example_btn = gr.Examples(
+                gr.Examples(
                     examples=examples,
                     inputs=prompt,
-                    label="Click on an example to try it",
+                    label="Click an example to try",
                     examples_per_page=5
                 )
-        
+
         # Footer
         gr.HTML(f"""
         <div class="footer">
             <p>Snowflake-G0-Release Demo • Created with Gradio • {datetime.datetime.now().year}</p>
         </div>
         """)
-        
-        # Set up interactions
+
+        # Interactions
         submit_btn.click(
             fn=generate_text,
             inputs=[prompt, temperature, top_p, top_k, max_new_tokens, history_state],
             outputs=[response_output, history_state, chat_history_display]
         )
-        
+
         prompt.submit(
             fn=generate_text,
             inputs=[prompt, temperature, top_p, top_k, max_new_tokens, history_state],
             outputs=[response_output, history_state, chat_history_display]
         )
-        
+
         clear_btn.click(
             fn=clear_conversation,
             inputs=[],
             outputs=[prompt, history_state, chat_history_display]
         )
-        
+
     return demo
 
-# Load model and tokenizer
-print("Loading Snowflake-G0-Release model and tokenizer...")
+# Initialize model
 try:
     model, tokenizer, pipeline = load_model_and_tokenizer()
-    print("Model loaded successfully!")
 except Exception as e:
     print(f"Error loading model: {str(e)}")
-    # Create a simple error demo if model fails to load
     with gr.Blocks(css=css) as error_demo:
         gr.HTML(f"""
         <div class="header" style="background-color: #ffebee;">
             <h1><span class="snowflake-icon">⚠️</span> Error Loading Model</h1>
-            <p>There was a problem loading the Snowflake-G0-Release model: {str(e)}</p>
+            <p>There was a problem loading the model: {str(e)}</p>
         </div>
         """)
     demo = error_demo
-
-# Create and launch the demo
-demo = create_demo()
+else:
+    demo = create_demo()
 
 # Launch the app
 if __name__ == "__main__":