Upload 3 files

model_Custm.py

@@ -158,6 +158,7 @@ class Wildnerve_tlm01(nn.Module, AbstractModel):
         super().__init__()
         # Set device once at the start
         object.__setattr__(self, "device", torch.device('cuda' if torch.cuda.is_available() else 'cpu'))
+        logger.info(f"Model initialized on device: {torch.device('cuda' if torch.cuda.is_available() else 'cpu')}")
         self.specialization = specialization
         self.dataset_path = dataset_path
         self.model_name = model_name

model_List.py

@@ -32,104 +32,64 @@ class PromptAnalyzer:
       - Provides candidate model identifiers or a single best match.
     """
     def __init__(self):
-        # Predefined topics with keyword sets for topic understanding
-        self.predefined_topics: Dict[str, List[str]] = {
-            "general": ["general", "overview", "basic", "introduction"],
-            "programming": ["code", "programming", "debug", "software", "algorithm", "bug"],
-            "science": ["research", "experiment", "science", "physics", "biology", "chemistry"],
-            "history": ["history", "ancient", "modern", "civilization", "war"],
-            "mathematics": ["math", "algebra", "calculus", "geometry", "statistics"]
+        self.logger = logging.getLogger(__name__)
+        
+        # Define topic keywords
+        self.predefined_topics = {
+            "programming": ["code", "function", "class", "algorithm", "programming", "python", "javascript", "java", "c++", "developer", "api"],
+            "science": ["science", "physics", "chemistry", "biology", "scientific", "experiment", "hypothesis", "theory"],
+            "mathematics": ["math", "equation", "calculus", "algebra", "geometry", "theorem", "mathematical"],
+            "history": ["history", "historical", "ancient", "century", "war", "civilization", "empire"]
         }
-        # Initialize a lightweight transformer encoder for embeddings
-        self.tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased")
-        self.encoder = AutoModel.from_pretrained("distilbert-base-uncased")
-        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-        self.encoder.to(self.device)
-        # Initialize SmartHybridAttention for refined representations
-        attention_config = get_hybrid_attention_config()
-        self.attention = SmartHybridAttention(attention_config)
-        self.attention.to(self.device)
-        logger.info("PromptAnalyzer initialized with DistilBERT and SmartHybridAttention.")
-
-    def _encode_text(self, text: str) -> np.ndarray:
-        """
-        Encode text into an embedding vector.
-        First, obtain token embeddings using DistilBERT.
-        Then refine these embeddings with SmartHybridAttention.
-        Finally, average-pool to produce a single vector.
-        """
-        inputs = self.tokenizer(text, return_tensors="pt", truncation=True, max_length=128)
-        inputs = {k: v.to(self.device) for k, v in inputs.items()}
-        with torch.no_grad():
-            outputs = self.encoder(**inputs)  # shape: [batch, seq_len, dim]
-        token_embeds = outputs.last_hidden_state  # [1, seq_len, dim]
-        # Transpose for attention: [seq_len, batch, dim]
-        token_embeds = token_embeds.transpose(0, 1)
-        attended, _ = self.attention(query=token_embeds, key=token_embeds, value=token_embeds)
-        # Transpose back and pool over tokens: [batch, seq_len, dim] -> [batch, dim]
-        attended = attended.transpose(0, 1)
-        pooled = attended.mean(dim=1)
-        return pooled.squeeze().cpu().numpy()
-
-    def analyze_prompt(self, prompt: str) -> Tuple[str, List[str]]:
-        """
-        Analyze the given prompt:
-          - Compute its refined embedding.
-          - For each predefined topic, encode its keyword string.
-          - Compute cosine similarity between prompt and topic embeddings.
-          - Return the primary topic (highest similarity) and any subtopics
-            with similarity above 80% of the top score.
-        """
-        prompt_embedding = self._encode_text(prompt)
+        
+        # IMPORTANT CHANGE: Don't load AutoModel, directly use model_Custm.Wildnerve_tlm01
+        try:
+            # Import the Wildnerve model directly - no AutoModel usage
+            from model_Custm import Wildnerve_tlm01
+            self.model_class = Wildnerve_tlm01
+            self.logger.info("Successfully imported Wildnerve_tlm01 from model_Custm")
+        except Exception as e:
+            self.logger.warning(f"Failed to import Wildnerve_tlm01: {e}")
+            self.model_class = None
+    
+    def analyze_prompt(self, prompt):
+        """Analyze prompt to determine primary and secondary topics"""
+        # Simple keyword-based classification
+        prompt_lower = prompt.lower()
         topic_scores = {}
+        
         for topic, keywords in self.predefined_topics.items():
-            topic_text = " ".join(keywords)
-            topic_embedding = self._encode_text(topic_text)
-            similarity = cosine_similarity(
-                prompt_embedding.reshape(1, -1),
-                topic_embedding.reshape(1, -1)
-            )[0][0]
-            topic_scores[topic] = similarity
-        sorted_topics = sorted(topic_scores.items(), key=lambda x: x[1], reverse=True)
-        primary_topic = sorted_topics[0][0] if sorted_topics else "general"
-        threshold = sorted_topics[0][1] * 0.8 if sorted_topics else 0.0
-        subtopics = [topic for topic, score in sorted_topics if score >= threshold and topic != primary_topic]
-        logger.debug(f"Prompt analyzed (first 30 chars): '{prompt[:30]}...' -> Primary: {primary_topic}, Subtopics: {subtopics}")
-        return primary_topic, subtopics
-
+            score = sum(1 for keyword in keywords if keyword in prompt_lower)
+            topic_scores[topic] = score
+        
+        # Find the topic with the highest score
+        if not topic_scores or max(topic_scores.values()) == 0:
+            return "general", []
+            
+        primary_topic = max(topic_scores.items(), key=lambda x: x[1])[0]
+        
+        # Get secondary topics (any with non-zero scores except primary)
+        secondary_topics = [t for t, s in topic_scores.items() 
+                            if s > 0 and t != primary_topic]
+        
+        return primary_topic, secondary_topics
+    
     def get_selected_models(self):
         """Return the list of selected models, always with model_Custm as primary."""
-        # Always prioritize model_Custm for all specializations
+        # Always use model_Custm.py as the primary model
         return ["model_Custm.py", "model_PrTr.py"]
         
     def choose_model(self, prompt=None):
-        """Choose model_Custm regardless of prompt content."""
+        """Always choose model_Custm regardless of prompt content"""
+        if self.model_class:
+            return self.model_class
+        
+        # Try importing again if initial import failed
         try:
-            # Ensure model_Custm is imported and registered
-            import importlib.util
-            import os
-            
-            # Get the directory containing this file
-            this_dir = os.path.dirname(os.path.abspath(__file__))
-            
-            # Load model_Custm
-            model_path = os.path.join(this_dir, "model_Custm.py")
-            if os.path.exists(model_path):
-                spec = importlib.util.spec_from_file_location("model_custm", model_path)
-                model_module = importlib.util.module_from_spec(spec)
-                spec.loader.exec_module(model_module)
-                
-                # Register in service registry
-                from service_registry import registry, MODEL, ensure_models_registered
-                ensure_models_registered()  # Make sure it's registered
-                
-                # Return the model class
-                return model_module.Wildnerve_tlm01
-            else:
-                self.logger.error(f"model_Custm.py not found at {model_path}")
-                return None
-        except Exception as e:
-            self.logger.error(f"Error in choose_model: {e}")
+            from model_Custm import Wildnerve_tlm01
+            return Wildnerve_tlm01
+        except ImportError as e:
+            self.logger.error(f"Failed to import Wildnerve_tlm01: {e}")
             return None
 
 # Register the PromptAnalyzer in the service registry to resolve dependencies.

transformer_patches.py

@@ -260,3 +260,47 @@ def apply_patch_to_layer(layer):
         return out
 
     layer.forward = forward_with_debug
+
+"""
+Patches for the transformers library to ensure compatibility
+"""
+import logging
+from types import FunctionType
+
+logger = logging.getLogger(__name__)
+
+def apply_transformers_patches():
+    """Apply patches to transformers library"""
+    try:
+        import torch
+        import transformers
+        
+        # Only apply safe patches that don't interfere with GPU usage
+        # Don't replace torch.device with a CPU-only version!
+        
+        # Fix AutoModel.from_pretrained to handle device mapping safely
+        if hasattr(transformers, 'AutoModel'):
+            original_from_pretrained = transformers.AutoModel.from_pretrained
+            
+            def safe_from_pretrained(*args, **kwargs):
+                # Keep any device_map parameter but handle it safely
+                if 'device_map' in kwargs and not isinstance(kwargs['device_map'], (str, dict)):
+                    logger.info("Fixing invalid device_map parameter")
+                    kwargs['device_map'] = "auto" if torch.cuda.is_available() else None
+                    
+                # Use cuda for faster performance if available
+                if 'torch_dtype' not in kwargs:
+                    kwargs['torch_dtype'] = torch.float16 if torch.cuda.is_available() else torch.float32
+                    
+                return original_from_pretrained(*args, **kwargs)
+            
+            transformers.AutoModel.from_pretrained = safe_from_pretrained
+            logger.info("Applied patch to AutoModel.from_pretrained that preserves GPU usage")
+        
+        return True
+    except Exception as e:
+        logger.error(f"Failed to apply transformers patches: {e}")
+        return False
+
+# Apply patches when module is imported
+apply_transformers_patches()