ss

app.py

@@ -9,18 +9,36 @@ import os
 import tempfile
 import base64
 from rank_bm25 import BM25Okapi
-from transformers import AutoModel, AutoTokenizer
-from sentence_transformers import SentenceTransformer
-from nltk.tokenize import word_tokenize, sent_tokenize
-from tqdm import tqdm
 import re
 import io
 import PyPDF2
 from docx import Document
 import csv
+
+# Fix for Replicate issue - must be before model imports
+import torch
+try:
+    from einops.layers.torch import Replicate
+except ImportError:
+    # Define our own Replicate class if not available
+    class Replicate(torch.nn.Module):
+        def __init__(self, *args, **kwargs):
+            super().__init__()
+        def forward(self, x):
+            return x
+
+# Make sure Replicate is available in the right module
+import sys
+import einops.layers.torch
+einops.layers.torch.Replicate = Replicate
+sys.modules['einops.layers.torch'].Replicate = Replicate
+
+# Now import models
+from transformers import AutoModel, AutoTokenizer
+from sentence_transformers import SentenceTransformer
+from nltk.tokenize import word_tokenize, sent_tokenize
+from tqdm import tqdm
 from explanation_generator import ExplanationGenerator
-from einops.layers.torch import Rearrange, Reduce
-from einops import rearrange, reduce, repeat
 
 # Download NLTK resources
 try:
@@ -33,21 +51,6 @@ EMBEDDING_MODEL_NAME = "nvidia/NV-Embed-v2"
 print(f"Loading embedding model {EMBEDDING_MODEL_NAME}...")
 
 try:
-    # Create a simple Replicate class since it's missing
-    class Replicate(torch.nn.Module):
-        def __init__(self):
-            super().__init__()
-        
-        def forward(self, x):
-            return x
-    
-    # Add to global namespace to make it available
-    import sys
-    sys.modules['__main__'].Replicate = Replicate
-    
-    # Also add to globals
-    globals()['Replicate'] = Replicate
-    
     # Load embedding model and tokenizer
     global_embedding_tokenizer = AutoTokenizer.from_pretrained(EMBEDDING_MODEL_NAME, trust_remote_code=True)
     global_embedding_model = AutoModel.from_pretrained(EMBEDDING_MODEL_NAME, trust_remote_code=True, device_map="auto")

explanation_generator.py

@@ -6,6 +6,23 @@ using the QwQ-32B model from Hugging Face.
 """
 
 import torch
+# Fix for Replicate issue - must be before model imports
+try:
+    from einops.layers.torch import Replicate
+except ImportError:
+    # Define our own Replicate class if not available
+    class Replicate(torch.nn.Module):
+        def __init__(self, *args, **kwargs):
+            super().__init__()
+        def forward(self, x):
+            return x
+
+# Make sure Replicate is available in the right module
+import sys
+import einops.layers.torch
+einops.layers.torch.Replicate = Replicate
+sys.modules['einops.layers.torch'].Replicate = Replicate
+
 from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
 import os
 import re
@@ -17,21 +34,6 @@ print("Loading Qwen/QwQ-32B model with 4-bit quantization...")
 QWQ_MODEL_NAME = "Qwen/QwQ-32B"
 
 try:
-    # Create a simple Replicate class since it's missing
-    class Replicate(torch.nn.Module):
-        def __init__(self):
-            super().__init__()
-        
-        def forward(self, x):
-            return x
-    
-    # Add to global namespace to make it available
-    import sys
-    sys.modules['__main__'].Replicate = Replicate
-    
-    # Also add to globals
-    globals()['Replicate'] = Replicate
-    
     # Configure 4-bit quantization for better performance
     quantization_config = BitsAndBytesConfig(
         load_in_4bit=True,