reverted back to old version

app/main.py

@@ -1,5 +1,5 @@
 import os
-from transformers import pipeline, BitsAndBytesConfig
+from transformers import pipeline
 import torch
 import nltk
 from nltk.translate.bleu_score import sentence_bleu, SmoothingFunction
@@ -418,10 +418,7 @@ async def startup_event():
     logger = logging.getLogger(__name__)
     logger.info("Starting application initialization...")
 
-    # Set PyTorch memory management settings
-    os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'expandable_segments:True'
-    
-    # Check if CUDA is available and has enough memory
+    # Check if CUDA is available
     device = "cuda" if torch.cuda.is_available() else "cpu"
     logger.info(f"Using device: {device}")
 
@@ -441,24 +438,15 @@ async def startup_event():
         logger.error(f"Error downloading NLTK data: {str(e)}")
         raise Exception(f"Failed to initialize application: {str(e)}")
 
-    # Initialize the model and index with quantization
+    # Initialize the model and index
     try:
-        # Configure 8-bit quantization
-        quantization_config = BitsAndBytesConfig(
-            load_in_8bit=True,
-            llm_int8_threshold=6.0,  # Default threshold for good balance of performance/accuracy
-            llm_int8_skip_modules=None,  # No modules to skip
-            llm_int8_enable_fp32_cpu_offload=True  # Enable CPU offloading if needed
-        )
-
         app.state.pipe = pipeline(
             "text-generation",
             model=MODEL_ID,
             trust_remote_code=True,
             token=HUGGINGFACE_TOKEN,
             device_map="auto",
-            torch_dtype=torch.float16 if device == "cuda" else torch.float32,
-            quantization_config=quantization_config  # Add quantization config
+            torch_dtype=torch.float16 if device == "cuda" else torch.float32
         )
         
         faiss_index, documents, embedding_model = await load_or_create_index()