Update app.py

app.py

@@ -70,13 +70,13 @@ class ModelManager:
             self.model = AutoModelForCausalLM.from_pretrained(
                 model_name,
                 token=HUGGINGFACE_TOKEN,
-                device_map="auto",  # Automatically handle device placement
-                torch_dtype=torch.float16,  # Use float16 to reduce memory usage
-                low_cpu_mem_usage=True,  # Optimize CPU memory usage
-                use_safetensors=True,  # Use safetensors for better memory management
-                max_memory={0: "6GB"},  # Limit GPU memory usage
-                offload_folder="offload",  # Folder for offloading to CPU
-                offload_state_dict=True  # Offload state dict to CPU
+                device_map="auto",
+                torch_dtype=torch.float16,
+                low_cpu_mem_usage=True,
+                use_safetensors=True,
+                max_memory={0: "6GB"},
+                offload_folder="offload",
+                offload_state_dict=True
             )
 
             # Create pipeline with minimal settings
@@ -86,8 +86,8 @@ class ModelManager:
                 "text-generation",
                 model=self.model,
                 tokenizer=self.tokenizer,
-                device_map="auto",  # Automatically handle device placement
-                torch_dtype=torch.float16,  # Use float16 for memory efficiency
+                device_map="auto",
+                torch_dtype=torch.float16,
                 max_new_tokens=512,
                 do_sample=True,
                 temperature=0.7,