Update app.py

app.py

@@ -6,7 +6,8 @@ from safetensors.torch import load_file  # Import safetensors for loading .safet
 import datetime
 
 # Model Constants
-MODEL_ID = "FlameF0X/Snowflake-G0-Release"  # HF repo when published
+MODEL_ID_V1 = "FlameF0X/Snowflake-G0-Release"
+MODEL_ID_V2 = "FlameF0X/Snowflake-G0-Release-2"
 MAX_LENGTH = 384
 TEMPERATURE_MIN = 0.1
 TEMPERATURE_MAX = 2.0
@@ -68,46 +69,84 @@ css = """
     margin-top: 0;
     color: #66ffaa;
 }
+.model-select {
+    background-color: #2a2a4a;
+    padding: 10px;
+    border-radius: 8px;
+    margin-bottom: 15px;
+}
 """
 
-# Helper functions to load model
-def load_model_and_tokenizer():
-    global model, tokenizer, pipeline  # Add this line
+# Global variables for models and tokenizers
+model_v1 = None
+tokenizer_v1 = None
+pipeline_v1 = None
+model_v2 = None
+tokenizer_v2 = None
+pipeline_v2 = None
 
-    # Load the tokenizer
-    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
+# Helper functions to load models
+def load_models_and_tokenizers():
+    global model_v1, tokenizer_v1, pipeline_v1, model_v2, tokenizer_v2, pipeline_v2
     
-    # Check if the pad_token is None, set it to eos_token if needed
-    if tokenizer.pad_token is None:
-        tokenizer.pad_token = tokenizer.eos_token
+    # Load the first model
+    print(f"Loading model from {MODEL_ID_V1}...")
+    tokenizer_v1 = AutoTokenizer.from_pretrained(MODEL_ID_V1)
+    if tokenizer_v1.pad_token is None:
+        tokenizer_v1.pad_token = tokenizer_v1.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_file_path = os.path.join(MODEL_ID_V1, "model.safetensors")
+    
+    if os.path.exists(model_file_path):
+        print("Loading model from safetensors file...")
+        model_v1 = load_file(model_file_path)
+    else:
+        print("Loading model from .bin file...")
+        model_v1 = AutoModelForCausalLM.from_pretrained(
+            MODEL_ID_V1, 
+            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+            device_map="auto"
+        )
+    
+    pipeline_v1 = TextGenerationPipeline(
+        model=model_v1,
+        tokenizer=tokenizer_v1,
+        return_full_text=False,
+        max_length=MAX_LENGTH
+    )
+    
+    # Load the second model
+    print(f"Loading model from {MODEL_ID_V2}...")
+    tokenizer_v2 = AutoTokenizer.from_pretrained(MODEL_ID_V2)
+    if tokenizer_v2.pad_token is None:
+        tokenizer_v2.pad_token = tokenizer_v2.eos_token
 
+    model_file_path = os.path.join(MODEL_ID_V2, "model.safetensors")
+    
     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
+        model_v2 = load_file(model_file_path)
     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")
+        model_v2 = AutoModelForCausalLM.from_pretrained(
+            MODEL_ID_V2, 
+            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,
+    pipeline_v2 = TextGenerationPipeline(
+        model=model_v2,
+        tokenizer=tokenizer_v2,
         return_full_text=False,
         max_length=MAX_LENGTH
     )
     
-    return model, tokenizer, pipeline
+    return (model_v1, tokenizer_v1, pipeline_v1), (model_v2, tokenizer_v2, pipeline_v2)
 
 # Helper functions for generation
 def generate_text(
-    prompt, 
+    prompt,
+    model_version, 
     temperature=TEMPERATURE_DEFAULT, 
     top_p=TOP_P_DEFAULT, 
     top_k=TOP_K_DEFAULT, 
@@ -121,6 +160,16 @@ def generate_text(
     history.append({"role": "user", "content": prompt})
     
     try:
+        # Select the appropriate pipeline based on model version
+        if model_version == "G0-Release":
+            pipeline = pipeline_v1
+            tokenizer = tokenizer_v1
+            model_name = "Snowflake-G0-Release"
+        else:  # "G0-Release-2"
+            pipeline = pipeline_v2
+            tokenizer = tokenizer_v2
+            model_name = "Snowflake-G0-Release-2"
+        
         # Generate response
         outputs = pipeline(
             prompt,
@@ -136,19 +185,23 @@ def generate_text(
         response = outputs[0]["generated_text"]
         
         # Add model response to history
-        history.append({"role": "assistant", "content": response})
+        history.append({"role": "assistant", "content": response, "model": model_name})
         
         # Format chat history for display
         formatted_history = []
         for entry in history:
-            role_prefix = "👤 User: " if entry["role"] == "user" else "❄️ Snowflake: "
+            if entry["role"] == "user":
+                role_prefix = "👤 User: "
+            else:
+                model_indicator = f"[{entry.get('model', 'Snowflake')}]"
+                role_prefix = f"❄️ {model_indicator}: "
             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}"})
+        history.append({"role": "assistant", "content": f"[ERROR] {error_msg}", "model": model_version})
         return error_msg, history, str(history)
 
 def clear_conversation():
@@ -172,36 +225,22 @@ def create_demo():
         # Header
         gr.HTML("""
         <div class="header">
-            <h1><span class="snowflake-icon">❄️</span> Snowflake-G0-Release Demo</h1>
-            <p>Experience the capabilities of the Snowflake-G0-Release language model</p>
+            <h1><span class="snowflake-icon">❄️</span> Snowflake Models Demo</h1>
+            <p>Experience the capabilities of the Snowflake series language models</p>
         </div>
         """)
         
-        # Model info
-        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.
-            
-            ### Model details
-            - Architecture: SnowflakeCore
-            - Hidden size: 384
-            - Number of attention heads: 6
-            - Number of layers: 4
-            - Feed-forward dimension: 768
-            - Maximum sequence 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
-            """)
-        
         # Chat interface
         with gr.Column():
+            # Model selection
+            with gr.Row(elem_classes="model-select"):
+                model_version = gr.Radio(
+                    ["G0-Release", "G0-Release-2"],
+                    label="Select Model Version",
+                    value="G0-Release-2",
+                    info="Choose which Snowflake model to use"
+                )
+            
             chat_history_display = gr.Textbox(
                 value="", 
                 label="Conversation History", 
@@ -288,20 +327,20 @@ def create_demo():
         # Footer
         gr.HTML(f"""
         <div class="footer">
-            <p>Snowflake-G0-Release Demo • Created with Gradio • {datetime.datetime.now().year}</p>
+            <p>Snowflake Models Demo • Created with Gradio • {datetime.datetime.now().year}</p>
         </div>
         """)
         
         # Set up interactions
         submit_btn.click(
             fn=generate_text,
-            inputs=[prompt, temperature, top_p, top_k, max_new_tokens, history_state],
+            inputs=[prompt, model_version, 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],
+            inputs=[prompt, model_version, temperature, top_p, top_k, max_new_tokens, history_state],
             outputs=[response_output, history_state, chat_history_display]
         )
         
@@ -313,19 +352,19 @@ def create_demo():
         
     return demo
 
-# Load model and tokenizer
-print("Loading Snowflake-G0-Release model and tokenizer...")
+# Load models and tokenizers
+print("Loading Snowflake models and tokenizers...")
 try:
-    model, tokenizer, pipeline = load_model_and_tokenizer()
-    print("Model loaded successfully!")
+    (model_v1, tokenizer_v1, pipeline_v1), (model_v2, tokenizer_v2, pipeline_v2) = load_models_and_tokenizers()
+    print("Models loaded successfully!")
 except Exception as e:
-    print(f"Error loading model: {str(e)}")
-    # Create a simple error demo if model fails to load
+    print(f"Error loading models: {str(e)}")
+    # Create a simple error demo if models fail 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>
+            <h1><span class="snowflake-icon">⚠️</span> Error Loading Models</h1>
+            <p>There was a problem loading the Snowflake models: {str(e)}</p>
         </div>
         """)
     demo = error_demo