Update app.py

app.py

@@ -43,14 +43,16 @@ model_name = "unsloth/Llama-3.2-3B-Instruct"
 # )
 
 model_config = transformers.AutoConfig.from_pretrained(model_name, max_new_tokens=1024)
+
+
 model = transformers.AutoModelForCausalLM.from_pretrained(
     model_name,
     trust_remote_code=True,
     config=model_config,
-    # quantization_config=bnb_config,
-    device_map=device,
+    device_map="auto" if device == "cuda" else None,
 )
 
+
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 
 query_pipeline = transformers.pipeline(
@@ -58,14 +60,15 @@ query_pipeline = transformers.pipeline(
     model=model,
     tokenizer=tokenizer,
     return_full_text=True,
-    torch_dtype=torch.float16,
-    device_map=device,
+    torch_dtype=torch.float16 if device == "cuda" else torch.float32,
+    device_map="auto" if device == "cuda" else None,
     temperature=0.7,
     top_p=0.9,
     top_k=50,
-    max_new_tokens=256
+    max_new_tokens=128  # Reduce this from 256
 )
 
+
 llm = HuggingFacePipeline(pipeline=query_pipeline)
 
 books_db_client_retriever = RetrievalQA.from_chain_type(