Upload app.py

app.py

@@ -29,14 +29,12 @@ def main():
     probe.load_state_dict(torch.load('probe.pt'))
 
     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()
 
     @torch.no_grad()
     def run_extraction(prompt):
         builder = SimplePairPromptBuilder(criterion='more positive')
         lst = [x.strip() for x in prompt.lower()[:300].split(',')][:100]
-        exp = runner.run_extraction(lst, lst, layers=[15], num_repeat=100, builder=builder, parallel=False, run_inference=True, debug=True, max_new_tokens=2)
+        exp = runner.run_extraction(lst, lst, layers=[15], num_repeat=50, builder=builder, parallel=False, run_inference=True, debug=True, max_new_tokens=2)
         test_ds = exp.make_dataset(15)
 
         import torch