Update app.py

Another attempt at reducing the RAM footprint without slowing down the model.

app.py

@@ -1,24 +1,30 @@
-# import gradio as gr
-# import torch
-# from transformers import AutoTokenizer, AutoModelForCausalLM
+import gradio as gr
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM
 
-# MODEL_ID = "HuggingFaceTB/SmolLM3-3B"
-# DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+MODEL_ID = "HuggingFaceTB/SmolLM3-3B"
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 
-# print("Loading tokenizer & model…")
-# tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
+print("Loading tokenizer & model…")
+tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
 # model = AutoModelForCausalLM.from_pretrained(MODEL_ID, torch_dtype=torch.bfloat16).to(DEVICE)
 
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL_ID,
+    load_in_8bit=True,  # or try load_in_4bit=True
+    device_map="cpu"
+)
+
 #########
 
-print("Loading tokenizer & model…")
-import gradio as gr
-from transformers import AutoTokenizer
-from optimum.onnxruntime import ORTModelForCausalLM
+# print("Loading tokenizer & model…")
+# import gradio as gr
+# from transformers import AutoTokenizer
+# from optimum.onnxruntime import ORTModelForCausalLM
 
-MODEL_ID = "HuggingFaceTB/SmolLM3-3B"
-tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
-model = ORTModelForCausalLM.from_pretrained(MODEL_ID, export=True, quantize=True)
+# MODEL_ID = "HuggingFaceTB/SmolLM3-3B"
+# tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
+# model = ORTModelForCausalLM.from_pretrained(MODEL_ID, export=True, quantize=True)
 
 #########