Update app.py

app.py

@@ -1,11 +1,16 @@
 import os
 import time
 import torch
+import warnings
 from flask import Flask, request, jsonify
 from flask_cors import CORS
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoModelForCausalLM, AutoTokenizer, logging
 import gradio as gr
 
+# Suppress warnings 
+warnings.filterwarnings("ignore")
+logging.set_verbosity_error()
+
 # Global variables
 MODEL_ID = "microsoft/bitnet-b1.58-2B-4T"
 MAX_LENGTH = 2048
@@ -27,21 +32,29 @@ def load_model_and_tokenizer():
     
     print(f"Loading model: {MODEL_ID}")
     
-    # Load tokenizer
-    tokenizer = AutoTokenizer.from_pretrained(
-        MODEL_ID,
-        use_fast=True,
-    )
-    
-    # Load model with optimizations for limited resources
-    model = AutoModelForCausalLM.from_pretrained(
-        MODEL_ID,
-        device_map="auto",
-        torch_dtype=torch.bfloat16,
-        load_in_4bit=True,
-    )
-    
-    print("Model and tokenizer loaded successfully!")
+    try:
+        # Load tokenizer
+        tokenizer = AutoTokenizer.from_pretrained(
+            MODEL_ID,
+            use_fast=True,
+            trust_remote_code=True  # Added to trust remote code
+        )
+        
+        # Load model with optimizations for limited resources
+        model = AutoModelForCausalLM.from_pretrained(
+            MODEL_ID,
+            device_map="auto",
+            torch_dtype=torch.bfloat16,
+            load_in_4bit=True,
+            trust_remote_code=True  # Added to trust remote code
+        )
+        
+        print("Model and tokenizer loaded successfully!")
+    except Exception as e:
+        import traceback
+        print(f"Error loading model: {str(e)}")
+        print(traceback.format_exc())
+        raise
 
 # Initialize Flask app
 app = Flask(__name__)
@@ -205,4 +218,8 @@ if __name__ == "__main__":
     
     # Create and launch Gradio interface
     demo = create_ui()
-    demo.launch(server_name="0.0.0.0", server_port=7860, share=True)
+    demo.launch(server_name="0.0.0.0", server_port=7860, share=True)
+    
+    # Flask won't reach here when Gradio is running
+    # If you want to run Flask separately:
+    # app.run(host='0.0.0.0', port=int(os.environ.get('PORT', 7860)))