update1 app

app.py

@@ -2,15 +2,26 @@ import gradio as gr
 from transformers import pipeline
 
 # Load the model and tokenizer
+# def load_model():
+#     # Load the NuminaMath-72B-CoT model
+#     pipe = pipeline(
+#         "text-generation",
+#         model="AI-MO/NuminaMath-72B-CoT",
+#         torch_dtype="auto",
+#         device_map="auto"  # Automatically map to available GPU/CPU
+#     )
+#     return pipe
+
+from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+
 def load_model():
-    # Load the NuminaMath-72B-CoT model
-    pipe = pipeline(
-        "text-generation",
-        model="AI-MO/NuminaMath-72B-CoT",
-        torch_dtype="auto",
-        device_map="auto"  # Automatically map to available GPU/CPU
+    tokenizer = AutoTokenizer.from_pretrained("AI-MO/NuminaMath-72B-CoT")
+    model = AutoModelForCausalLM.from_pretrained(
+        "AI-MO/NuminaMath-72B-CoT",
+        device_map="auto",  # Automatically map to available GPU
+        offload_folder="offload"  # Offload unused parts to disk
     )
-    return pipe
+    return pipeline("text-generation", model=model, tokenizer=tokenizer)
 
 # Initialize the pipeline
 model_pipeline = load_model()