Update lib/graph_extract.py

lib/graph_extract.py

@@ -34,7 +34,7 @@ print("Initializing application...")
 model = AutoModelForCausalLM.from_pretrained(
     "sciphi/triplex", 
     trust_remote_code=True,
-    attn_implementation="flash_attention_2" if flash_attn_installed else None,
+    # attn_implementation="flash_attention_2" if flash_attn_installed else None,
     torch_dtype=torch.bfloat16,
     device_map="auto", 
     low_cpu_mem_usage=True,#advised if any device map given
@@ -43,7 +43,7 @@ model = AutoModelForCausalLM.from_pretrained(
 tokenizer = AutoTokenizer.from_pretrained(
     "sciphi/triplex", 
     trust_remote_code=True,
-    attn_implementation="flash_attention_2" if flash_attn_installed else None,
+    # attn_implementation="flash_attention_2" if flash_attn_installed else None,
         torch_dtype=torch.bfloat16,
     )