Update app.py

app.py

@@ -1,25 +1,39 @@
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
 from peft import PeftModel
+import gradio as gr
 
-# Load the tokenizer and base model
+# Set model name and path
 model_name = "ybelkada/falcon-7b-sharded-bf16"
 fine_tuned_model = "mounseflit/falcon-7b-marrakech"
 
+# Load tokenizer
 tokenizer = AutoTokenizer.from_pretrained(model_name)
-base_model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto")
 
-# Load the fine-tuned LoRA model
+# Load base model with 8-bit precision and offload to CPU
+base_model = AutoModelForCausalLM.from_pretrained(
+    model_name,
+    load_in_8bit=True,        # Quantization to 8-bit
+    device_map="auto",        # Auto device map for offloading
+    offload_folder="offload", # Offload large parts of the model to disk
+    offload_state_dict=True   # Enable state dict offloading to reduce memory usage
+)
+
+# Load the fine-tuned LoRA model on top of the quantized model
 model = PeftModel.from_pretrained(base_model, fine_tuned_model)
 
-# Define the function for generating text
+# Ensure the model is in evaluation mode
+model.eval()
+
+# Function to generate text
 def generate_text(prompt):
-    inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
-    outputs = model.generate(**inputs, max_length=200)
+    inputs = tokenizer(prompt, return_tensors="pt", max_length=50, truncation=True).to("cpu")  # Reduce input length
+    with torch.no_grad():
+        outputs = model.generate(**inputs, max_length=100)  # Reduce output length
     return tokenizer.decode(outputs[0], skip_special_tokens=True)
 
-# Gradio Interface
-import gradio as gr
+# Create Gradio interface
+iface = gr.Interface(fn=generate_text, inputs="text", outputs="text", title="Falcon 7B Lite")
 
-iface = gr.Interface(fn=generate_text, inputs="text", outputs="text")
+# Launch the app
 iface.launch()