Upload 2 files

app.py

@@ -28,7 +28,7 @@ def main():
     probe = BinaryProbe(config).cuda()
     probe.load_state_dict(torch.load('probe.pt'))
 
-    runner = PairwiseExtractionRunner.from_pretrained('mistralai/Mistral-7B-Instruct-v0.1', optimize=True, max_memory=mdict, device_map=dmap, low_cpu_mem_usage=True)
+    runner = PairwiseExtractionRunner.from_pretrained('mistralai/Mistral-7B-Instruct-v0.1', optimize=False, torch_dtype=torch.float16, max_memory=mdict, device_map=dmap, low_cpu_mem_usage=True)
     device = "cpu" if args.no_cuda else "cuda"
     lock = Lock()
 

requirements.txt

@@ -2,5 +2,4 @@ torch
 torchvision
 torchaudio
 gradio==3.36.1
-biasprobe
-flash-attn
+biasprobe