update

src/attribution/attntrace.py

@@ -6,6 +6,7 @@ import torch.nn.functional as F
 import gc
 from src.prompts import wrap_prompt_attention
 from .attention_utils import *
+import spaces
 
 class AttnTraceAttribution(Attribution):
     def __init__(self,  llm,explanation_level = "segment",K=5, avg_k=5, q=0.4, B=30, verbose =1):

src/models/Llama.py

@@ -4,6 +4,7 @@ from .Model import Model
 import os
 import signal
 from functools import lru_cache
+import spaces
 
 def handle_timeout(sig, frame):
     raise TimeoutError('took too long')