Update app.py

app.py

@@ -1,310 +1,53 @@
 import os
-import torch
-import gradio as gr
-from transformers import AutoTokenizer, AutoModelForCausalLM, TextGenerationPipeline
-import datetime
-
-# Model Constants
-MODEL_ID = "./model"  # Local folder containing model files
-MAX_LENGTH = 384
-TEMPERATURE_MIN = 0.1
-TEMPERATURE_MAX = 2.0
-TEMPERATURE_DEFAULT = 0.7
-TOP_P_MIN = 0.1
-TOP_P_MAX = 1.0
-TOP_P_DEFAULT = 0.9
-TOP_K_MIN = 1
-TOP_K_MAX = 100
-TOP_K_DEFAULT = 40
-MAX_NEW_TOKENS_MIN = 16
-MAX_NEW_TOKENS_MAX = 1024
-MAX_NEW_TOKENS_DEFAULT = 256
-
-# CSS for the app
-css = """
-.gradio-container {
-    background-color: #1e1e2f !important;
-    color: #e0e0e0 !important;
-}
-.header {
-    background-color: #2b2b3c;
-    padding: 20px;
-    margin-bottom: 20px;
-    border-radius: 10px;
-    text-align: center;
-}
-.header h1 {
-    color: #66ccff;
-    margin-bottom: 10px;
-}
-.snowflake-icon {
-    font-size: 24px;
-    margin-right: 10px;
-}
-.footer {
-    text-align: center;
-    margin-top: 20px;
-    font-size: 0.9em;
-    color: #999;
-}
-.parameter-section {
-    background-color: #2a2a3a;
-    padding: 15px;
-    border-radius: 8px;
-    margin-bottom: 15px;
-}
-.parameter-section h3 {
-    margin-top: 0;
-    color: #66ccff;
-}
-.example-section {
-    background-color: #223344;
-    padding: 15px;
-    border-radius: 8px;
-    margin-bottom: 15px;
-}
-.example-section h3 {
-    margin-top: 0;
-    color: #66ffaa;
-}
-"""
-
-# Helper function to load model and tokenizer
-def load_model_and_tokenizer():
-    global model, tokenizer, pipeline
-
-    print("Loading Snowflake-G0-Release model and tokenizer...")
-    
-    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
-
-    if tokenizer.pad_token is None:
-        tokenizer.pad_token = tokenizer.eos_token
-
-    model = AutoModelForCausalLM.from_pretrained(
-        MODEL_ID,
-        torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
-        device_map="auto"
-    )
-
-    pipeline = TextGenerationPipeline(
-        model=model,
-        tokenizer=tokenizer,
-        return_full_text=False,
-        max_length=MAX_LENGTH
-    )
-    
-    print("Model loaded successfully!")
-    return model, tokenizer, pipeline
-
-# Helper functions for generation
-def generate_text(
-    prompt, 
-    temperature=TEMPERATURE_DEFAULT, 
-    top_p=TOP_P_DEFAULT, 
-    top_k=TOP_K_DEFAULT, 
-    max_new_tokens=MAX_NEW_TOKENS_DEFAULT,
-    history=None
-):
-    if history is None:
-        history = []
-    
-    history.append({"role": "user", "content": prompt})
-
-    try:
-        outputs = pipeline(
-            prompt,
-            do_sample=temperature > 0,
-            temperature=temperature,
-            top_p=top_p,
-            top_k=top_k,
-            max_new_tokens=max_new_tokens,
-            pad_token_id=tokenizer.pad_token_id,
-            num_return_sequences=1
-        )
-
-        response = outputs[0]["generated_text"]
-        history.append({"role": "assistant", "content": response})
-
-        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}"})
-        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.",
-    "Explain the concept of artificial neural networks to a 10-year-old.",
-    "What are some interesting applications of natural language processing?",
-    "Write a haiku about programming.",
-    "Create a dialogue between two AI researchers discussing the future of language models."
-]
-
-# Main app creation
-def create_demo():
-    with gr.Blocks(css=css) as 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>
-        </div>
-        """)
-        
-        # About accordion
-        with gr.Accordion("About Snowflake-G0-Release", open=False):
-            gr.Markdown("""
-            ## Snowflake-G0-Release
-            
-            Initial release of the Snowflake series trained on DialogMLM-50K.
-            
-            ### Model Details
-            - Architecture: SnowflakeCore
-            - Hidden size: 384
-            - Attention heads: 6
-            - Layers: 4
-            - Feed-forward dim: 768
-            - Max seq length: 384
-            - Vocabulary size: 30522 (BERT tokenizer)
-            
-            ### 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(
-                value="", 
-                label="Conversation History", 
-                lines=10, 
-                max_lines=30,
-                interactive=False
-            )
-            history_state = gr.State([])
-
-            with gr.Row():
-                with gr.Column(scale=4):
-                    prompt = gr.Textbox(
-                        placeholder="Type your message here...", 
-                        label="Your Input",
-                        lines=2
-                    )
-                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", 
-                lines=5,
-                max_lines=10,
-                interactive=False
-            )
-
-        # Generation Parameters
-        with gr.Accordion("Generation Parameters", open=False):
-            with gr.Column(elem_classes="parameter-section"):
-                with gr.Row():
-                    with gr.Column():
-                        temperature = gr.Slider(
-                            minimum=TEMPERATURE_MIN, 
-                            maximum=TEMPERATURE_MAX,
-                            value=TEMPERATURE_DEFAULT,
-                            step=0.05,
-                            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)"
-                        )
-                    with gr.Column():
-                        top_k = gr.Slider(
-                            minimum=TOP_K_MIN,
-                            maximum=TOP_K_MAX,
-                            value=TOP_K_DEFAULT,
-                            step=1,
-                            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"
-                        )
+import json
+from safetensors.torch import load_file, save_file
+
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+# Path to your model folder
+model_dir = "./model"
+
+# Step 1: Fix config.json if missing model_type
+config_path = os.path.join(model_dir, "config.json")
+if os.path.exists(config_path):
+    with open(config_path, "r") as f:
+        config = json.load(f)
+
+    if "model_type" not in config:
+        print("⚙️ Adding missing 'model_type' to config.json...")
+        # You can adjust 'gpt2' to whatever your real model type is
+        config["model_type"] = "gpt2"
+        with open(config_path, "w") as f:
+            json.dump(config, f, indent=2)
+    else:
+        print("✅ 'model_type' already exists in config.json.")
+else:
+    raise FileNotFoundError("config.json not found in model directory!")
 
-        # Example Prompts
-        with gr.Accordion("Example Prompts", open=True):
-            with gr.Column(elem_classes="example-section"):
-                gr.Examples(
-                    examples=examples,
-                    inputs=prompt,
-                    label="Click an example to try",
-                    examples_per_page=5
-                )
+# Step 2: Fix .safetensors file metadata
+safetensors_files = [f for f in os.listdir(model_dir) if f.endswith(".safetensors")]
+if safetensors_files:
+    safetensors_path = os.path.join(model_dir, safetensors_files[0])
+    print(f"🛠 Fixing metadata in: {safetensors_path}")
 
-        # Footer
-        gr.HTML(f"""
-        <div class="footer">
-            <p>Snowflake-G0-Release Demo • Created with Gradio • {datetime.datetime.now().year}</p>
-        </div>
-        """)
+    state_dict = load_file(safetensors_path)
+    fixed_path = os.path.join(model_dir, "model_fixed.safetensors")
+    save_file(state_dict, fixed_path, metadata={"format": "pt"})
 
-        # 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]
-        )
+    print(f"✅ Saved fixed safetensors: {fixed_path}")
+else:
+    print("⚠️ No .safetensors file found to fix.")
 
-        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]
-        )
+# Step 3: Load model to verify it works now
+print("🚀 Trying to load the model...")
 
-        clear_btn.click(
-            fn=clear_conversation,
-            inputs=[],
-            outputs=[prompt, history_state, chat_history_display]
-        )
+model = AutoModelForCausalLM.from_pretrained(model_dir, trust_remote_code=True)
+tokenizer = AutoTokenizer.from_pretrained(model_dir)
 
-    return demo
+print("🎉 Model loaded successfully!")
 
-# Initialize model
-try:
-    model, tokenizer, pipeline = load_model_and_tokenizer()
-except Exception as e:
-    print(f"Error loading model: {str(e)}")
-    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 model: {str(e)}</p>
-        </div>
-        """)
-    demo = error_demo
-else:
-    demo = create_demo()
+# Step 4 (optional): Save model again safely
+model.save_pretrained(model_dir, safe_serialization=True)
+tokenizer.save_pretrained(model_dir)
 
-# Launch the app
-if __name__ == "__main__":
-    demo.launch()
+print("✅ Model and tokenizer saved safely with correct format!")