Update app.py

app.py

@@ -1,24 +1,11 @@
-#!/usr/bin/env python3
-# Hugging Face Space app.py for Snowflake-G0-Release demo
-
 import os
-import sys
 import gradio as gr
 import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM, TextGenerationPipeline
 import datetime
-import logging
-
-# Configure logging
-logging.basicConfig(
-    level=logging.INFO,
-    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
-    handlers=[logging.StreamHandler(sys.stdout)]
-)
-logger = logging.getLogger("snowflake-demo")
 
 # Model Constants
-MODEL_ID = os.environ.get("MODEL_ID", "Snowflake-G0-Release")  # Use environment variable or default
+MODEL_ID = "FlameF0X/Snowflake-G0-Release"  # Replace with actual HF repo when published
 MAX_LENGTH = 384
 TEMPERATURE_MIN = 0.1
 TEMPERATURE_MAX = 2.0
@@ -30,7 +17,7 @@ TOP_K_MIN = 1
 TOP_K_MAX = 100
 TOP_K_DEFAULT = 40
 MAX_NEW_TOKENS_MIN = 16
-MAX_NEW_TOKENS_MAX = 512
+MAX_NEW_TOKENS_MAX = 1024
 MAX_NEW_TOKENS_DEFAULT = 256
 
 # CSS for the app
@@ -81,44 +68,29 @@ css = """
 }
 """
 
-# Global variables for model, tokenizer, and pipeline
-model = None
-tokenizer = None
-pipeline = None
-
 # Helper functions
 def load_model_and_tokenizer():
-    global model, tokenizer, pipeline
-    
-    logger.info(f"Loading model and tokenizer from: {MODEL_ID}")
-    
     # Load tokenizer
     tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
     if tokenizer.pad_token is None:
         tokenizer.pad_token = tokenizer.eos_token
     
-    # Determine device and precision
-    device_map = "auto"
-    dtype = torch.float16 if torch.cuda.is_available() else torch.float32
-    logger.info(f"Using device: {device_map}, dtype: {dtype}")
-    
     # Load model with optimizations
     model = AutoModelForCausalLM.from_pretrained(
         MODEL_ID,
-        torch_dtype=dtype,
-        device_map=device_map,
-        low_cpu_mem_usage=True
+        torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+        device_map="auto"
     )
     
     # Create pipeline
     pipeline = TextGenerationPipeline(
         model=model,
         tokenizer=tokenizer,
-        return_full_text=False
+        return_full_text=False,
+        max_length=MAX_LENGTH
     )
     
-    logger.info("Model loaded successfully!")
-    return True
+    return model, tokenizer, pipeline
 
 def generate_text(
     prompt, 
@@ -131,16 +103,10 @@ def generate_text(
     if history is None:
         history = []
     
-    if not prompt.strip():
-        return "", history, "\n\n".join([f"{'👤 User: ' if h['role'] == 'user' else '❄️ Snowflake: '}{h['content']}" for h in history])
-    
     # Add current prompt to history
     history.append({"role": "user", "content": prompt})
     
     try:
-        logger.info(f"Generating text for prompt: {prompt[:50]}...")
-        logger.info(f"Parameters: temp={temperature}, top_p={top_p}, top_k={top_k}, max_tokens={max_new_tokens}")
-        
         # Generate response
         outputs = pipeline(
             prompt,
@@ -154,7 +120,6 @@ def generate_text(
         )
         
         response = outputs[0]["generated_text"]
-        logger.info(f"Generated response: {response[:50]}...")
         
         # Add model response to history
         history.append({"role": "assistant", "content": response})
@@ -165,23 +130,19 @@ def generate_text(
             role_prefix = "👤 User: " if entry["role"] == "user" else "❄️ Snowflake: "
             formatted_history.append(f"{role_prefix}{entry['content']}")
         
-        return "", history, "\n\n".join(formatted_history)
+        return response, history, "\n\n".join(formatted_history)
     
     except Exception as e:
-        logger.error(f"Error generating response: {str(e)}", exc_info=True)
         error_msg = f"Error generating response: {str(e)}"
         history.append({"role": "assistant", "content": f"[ERROR] {error_msg}"})
-        
-        formatted_history = []
-        for entry in history:
-            role_prefix = "👤 User: " if entry["role"] == "user" else "❄️ Snowflake: "
-            formatted_history.append(f"{role_prefix}{entry['content']}")
-        
-        return "", history, "\n\n".join(formatted_history)
+        return error_msg, history, str(history)
 
 def clear_conversation():
     return "", [], ""
 
+def apply_preset_example(example, history):
+    return example, history
+
 # Example prompts
 examples = [
     "Write a short story about a snowflake that comes to life.",
@@ -193,7 +154,7 @@ examples = [
 
 # Main function
 def create_demo():
-    with gr.Blocks(css=css, title="Snowflake-G0-Release Demo") as demo:
+    with gr.Blocks(css=css) as demo:
         # Header
         gr.HTML("""
         <div class="header">
@@ -229,8 +190,8 @@ def create_demo():
         with gr.Column():
             chat_history_display = gr.Textbox(
                 value="", 
-                label="Conversation", 
-                lines=12, 
+                label="Conversation History", 
+                lines=10, 
                 max_lines=30,
                 interactive=False
             )
@@ -247,9 +208,16 @@ def create_demo():
                         lines=2
                     )
                 with gr.Column(scale=1):
-                    with gr.Row():
-                        submit_btn = gr.Button("Send", variant="primary")
-                        clear_btn = gr.Button("Clear")
+                    submit_btn = gr.Button("Send", variant="primary")
+                    clear_btn = gr.Button("Clear Conversation")
+            
+            response_output = gr.Textbox(
+                value="", 
+                label="Model Response", 
+                lines=5,
+                max_lines=10,
+                interactive=False
+            )
         
         # Advanced parameters
         with gr.Accordion("Generation Parameters", open=False):
@@ -299,7 +267,8 @@ def create_demo():
                 example_btn = gr.Examples(
                     examples=examples,
                     inputs=prompt,
-                    label="Click on an example to try it"
+                    label="Click on an example to try it",
+                    examples_per_page=5
                 )
         
         # Footer
@@ -313,13 +282,13 @@ def create_demo():
         submit_btn.click(
             fn=generate_text,
             inputs=[prompt, temperature, top_p, top_k, max_new_tokens, history_state],
-            outputs=[prompt, history_state, chat_history_display]
+            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=[prompt, history_state, chat_history_display]
+            outputs=[response_output, history_state, chat_history_display]
         )
         
         clear_btn.click(
@@ -330,13 +299,13 @@ def create_demo():
         
     return demo
 
-# Load model and create demo
+# Load model and tokenizer
+print("Loading Snowflake-G0-Release model and tokenizer...")
 try:
-    logger.info("Starting Snowflake-G0-Release demo...")
-    load_model_and_tokenizer()
-    demo = create_demo()
+    model, tokenizer, pipeline = load_model_and_tokenizer()
+    print("Model loaded successfully!")
 except Exception as e:
-    logger.error(f"Failed to load model: {str(e)}", exc_info=True)
+    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"""
@@ -347,6 +316,9 @@ except Exception as e:
         """)
     demo = error_demo
 
+# Create and launch the demo
+demo = create_demo()
+
 # Launch the app
 if __name__ == "__main__":
     demo.launch()