fiix

app.py

@@ -24,80 +24,11 @@ class MultitaskCodeSimilarityModel(nn.Module):
         self.encoder = AutoModel.from_pretrained(model_name, config=self.config)
         self.classifier = nn.Linear(self.config.hidden_size, num_labels)
 
-        # For explanation generation
-        self.decoder_embedding = nn.Linear(self.config.hidden_size, self.config.hidden_size)
-        self.decoder = nn.GRU(
-            input_size=self.config.hidden_size,
-            hidden_size=self.config.hidden_size,
-            batch_first=True
-        )
-        self.explanation_head = nn.Linear(self.config.hidden_size, len(tokenizer))
-
-    def forward(self, input_ids, attention_mask, explanation_ids=None, explanation_mask=None):
+    def forward(self, input_ids, attention_mask):
         outputs = self.encoder(input_ids=input_ids, attention_mask=attention_mask)
         pooled = outputs.last_hidden_state[:, 0]
         logits = self.classifier(pooled)
-
-        explanation_logits = None
-        if explanation_ids is not None:
-            batch_size = input_ids.size(0)
-            seq_length = explanation_ids.size(1)
-
-            # Initialize decoder with pooled representation
-            decoder_input = self.decoder_embedding(pooled).unsqueeze(1).expand(-1, seq_length, -1)
-
-            # Run decoder
-            decoder_outputs, _ = self.decoder(decoder_input)
-
-            # Generate logits for each position
-            explanation_logits = self.explanation_head(decoder_outputs)
-
-        return logits, explanation_logits
-    
-    def generate_explanation(self, input_ids, attention_mask, max_length=128):
-        """Generate explanation text for inference"""
-        device = input_ids.device
-        
-        # Get encoding
-        outputs = self.encoder(input_ids=input_ids, attention_mask=attention_mask)
-        pooled = outputs.last_hidden_state[:, 0]
-        
-        # First token (usually [CLS] or <s>)
-        bos_token_id = self.tokenizer.bos_token_id if self.tokenizer.bos_token_id is not None else self.tokenizer.cls_token_id
-        current_token_id = torch.full((pooled.size(0), 1), bos_token_id, dtype=torch.long, device=device)
-        
-        generated_ids = [current_token_id]
-        
-        # Initial hidden state
-        hidden = pooled.unsqueeze(0)  # Add seq dimension for GRU
-        
-        for _ in range(max_length - 1):
-            # Get decoder input from current token
-            decoder_input = self.decoder_embedding(pooled).unsqueeze(1)
-            
-            # Run decoder one step
-            decoder_output, hidden = self.decoder(decoder_input, hidden)
-            
-            # Get next token probabilities
-            next_token_logits = self.explanation_head(decoder_output.squeeze(1))
-            next_token_id = torch.argmax(next_token_logits, dim=-1, keepdim=True)
-            
-            # Stop if we predict EOS
-            if (next_token_id == self.tokenizer.eos_token_id).all():
-                break
-                
-            generated_ids.append(next_token_id)
-            
-        # Concatenate all generated tokens
-        all_tokens = torch.cat(generated_ids, dim=1)
-        
-        # Convert to text
-        explanations = []
-        for tokens in all_tokens:
-            explanation = self.tokenizer.decode(tokens, skip_special_tokens=True)
-            explanations.append(explanation)
-            
-        return explanations
+        return logits
 
 # Load model and tokenizer
 try:
@@ -188,24 +119,16 @@ async def predict(data: SimilarityInput):
         # Get prediction from the model
         with torch.no_grad():
             # Check if using custom model or fallback
-            if hasattr(model, 'generate_explanation'):
+            if hasattr(model, 'forward'):
                 # Our custom model
-                logits, _ = model(
+                logits = model(
                     input_ids=inputs["input_ids"],
                     attention_mask=inputs["attention_mask"]
                 )
-                
-                # Generate explanation
-                explanations = model.generate_explanation(
-                    input_ids=inputs["input_ids"],
-                    attention_mask=inputs["attention_mask"]
-                )
-                explanation = explanations[0] if explanations else ""
             else:
                 # Fallback to standard model
                 outputs = model(**inputs)
                 logits = outputs.logits
-                explanation = ""
 
         # Process results
         probs = torch.softmax(logits, dim=-1)[0].cpu().tolist()
@@ -214,16 +137,6 @@ async def predict(data: SimilarityInput):
         # Map prediction to class name
         classification = label_to_class.get(prediction, "Unknown")
         
-        # Generate explanations contextually if not available from model
-        if not explanation or explanation.strip() == "":
-            # Template explanations based on classification
-            if classification == "Duplicate":
-                explanation = f"Tests {data.test_case_1.name} and {data.test_case_2.name} are duplicates because they both check the output formatting of their respective methods using the same approach of redirecting stdout to a buffer and verifying the exact output string."
-            elif classification == "Redundant":
-                explanation = f"Tests {data.test_case_1.name} and {data.test_case_2.name} are redundant because they test similar functionality (output formatting) using the same testing technique (capturing stdout) but on different methods."
-            elif classification == "Distinct":
-                explanation = f"Tests {data.test_case_1.name} and {data.test_case_2.name} are distinct because they test completely different functionality of the BankApp class: one tests listClients() while the other tests deposit()."
-
         return {
             "pair_id": data.pair_id,
             "test_case_1_name": data.test_case_1.name,
@@ -231,7 +144,6 @@ async def predict(data: SimilarityInput):
             "similarity": {
                 "score": prediction,
                 "classification": classification,
-                "explanation": explanation
             },
             "probabilities": probs
         }