add offload

app.py

@@ -2,7 +2,6 @@ import gradio as gr
 from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 from peft import PeftModel
 import torch
-
 import os
 
 # Pastikan Hugging Face Token disediakan (jika private repo)
@@ -12,11 +11,20 @@ hf_token = os.getenv('HF_TOKEN')
 base_model = "google/gemma-2b-it"
 adapter_model = "FadQ/gemma-2b-diary-consultaton-chatbot"
 
-# Load model dasar
-model = AutoModelForCausalLM.from_pretrained(base_model, torch_dtype=torch.float16, device_map="auto")
+# Load model dasar dengan offloading
+model = AutoModelForCausalLM.from_pretrained(
+    base_model, 
+    torch_dtype=torch.float16, 
+    device_map="auto", 
+    offload_folder="offload"  # Tambahkan folder untuk offloading ke disk
+)
 
-# Load adapter PEFT
-model = PeftModel.from_pretrained(model, adapter_model)
+# Load adapter PEFT dengan offloading
+model = PeftModel.from_pretrained(
+    model, 
+    adapter_model, 
+    offload_folder="offload"
+)
 
 # Load tokenizer
 tokenizer = AutoTokenizer.from_pretrained(base_model)
@@ -26,7 +34,8 @@ pipe = pipeline("text-generation", model=model, tokenizer=tokenizer, device=0)
 
 def predict(input_text):
     inputs = tokenizer(input_text, return_tensors="pt").to("cuda")
-    output = model.generate(**inputs, max_length=150)
+    with torch.no_grad():  # Hindari penggunaan memori yang tidak perlu
+        output = model.generate(**inputs, max_length=150)
     return tokenizer.decode(output[0], skip_special_tokens=True)
 
 # Create Gradio interface