Fixed indentation syntax error

app.py

@@ -64,114 +64,115 @@ def generate_step_by_step(history):
         
         prompt = history[-1][0]
         print(f"Generating for prompt: {prompt}")
-    
-    steps = 20
-    temp = 0.2
-    top_k = 10
-
-    bos_id = tokenizer.token_to_id("[BOS]")
-    eos_id = tokenizer.token_to_id("[EOS]")
-    mask_id = tokenizer.token_to_id("[MASK]")
-    pad_id = tokenizer.token_to_id("[PAD]")
-    unk_id = tokenizer.token_to_id("[UNK]")
-    special_ids = {bos_id, eos_id, mask_id, pad_id, unk_id}
+        
+        steps = 20
+        temp = 0.2
+        top_k = 10
 
-    formatted_prompt = f"user: {prompt.lower().strip()} bot:"
-    input_ids = tokenizer.encode(formatted_prompt).ids
+        bos_id = tokenizer.token_to_id("[BOS]")
+        eos_id = tokenizer.token_to_id("[EOS]")
+        mask_id = tokenizer.token_to_id("[MASK]")
+        pad_id = tokenizer.token_to_id("[PAD]")
+        unk_id = tokenizer.token_to_id("[UNK]")
+        special_ids = {bos_id, eos_id, mask_id, pad_id, unk_id}
 
-    max_resp = min(40, MAX_SEQ_LENGTH - len(input_ids) - 2)
+        formatted_prompt = f"user: {prompt.lower().strip()} bot:"
+        input_ids = tokenizer.encode(formatted_prompt).ids
 
-    sequence = [bos_id] + input_ids + [mask_id] * max_resp + [eos_id]
-    sequence += [pad_id] * (MAX_SEQ_LENGTH - len(sequence))
-    seq_tensor = torch.tensor([sequence], dtype=torch.long, device=device)
+        max_resp = min(40, MAX_SEQ_LENGTH - len(input_ids) - 2)
 
-    response_start = 1 + len(input_ids)
-    response_end = response_start + max_resp
-    mask_indices = list(range(response_start, response_end))
-    num_masks = len(mask_indices)
+        sequence = [bos_id] + input_ids + [mask_id] * max_resp + [eos_id]
+        sequence += [pad_id] * (MAX_SEQ_LENGTH - len(sequence))
+        seq_tensor = torch.tensor([sequence], dtype=torch.long, device=device)
 
-    running_confidence = torch.zeros(num_masks, device=device)
+        response_start = 1 + len(input_ids)
+        response_end = response_start + max_resp
+        mask_indices = list(range(response_start, response_end))
+        num_masks = len(mask_indices)
 
-    def get_output(seq_tensor, is_final=False):
-        response_ids = seq_tensor[0][response_start:response_end]
-        return decode_with_masks(response_ids, is_final)
+        running_confidence = torch.zeros(num_masks, device=device)
 
-    out_text = get_output(seq_tensor, is_final=False)
-    history[-1][1] = out_text
-    yield history
+        def get_output(seq_tensor, is_final=False):
+            response_ids = seq_tensor[0][response_start:response_end]
+            return decode_with_masks(response_ids, is_final)
 
-    current_seq = seq_tensor.squeeze(0).clone()
+        out_text = get_output(seq_tensor, is_final=False)
+        history[-1][1] = out_text
+        yield history
 
-    # Pre-calculate penalty indices once
-    vocab_size = model.fc_out.out_features
-    penalty_indices = torch.tensor([i for i in range(vocab_size) if i not in special_ids], device=device)
+        current_seq = seq_tensor.squeeze(0).clone()
 
-    for step in range(1, steps + 1):
-        t_val = max(1.0 - step / steps, 0.05)
-        t = torch.tensor([t_val], device=device)
-        pad_mask = (seq_tensor == pad_id)
+        # Pre-calculate penalty indices once
+        vocab_size = model.fc_out.out_features
         
-        with torch.no_grad():
-            logits = model(seq_tensor, t, src_key_padding_mask=pad_mask)
+        for step in range(1, steps + 1):
+            t_val = max(1.0 - step / steps, 0.05)
+            t = torch.tensor([t_val], device=device)
+            pad_mask = (seq_tensor == pad_id)
             
-        response_logits = logits[0, mask_indices]
-        
-        # Optimized Vectorized Repetition Penalty
-        if step > 1:
-            prev_pred = current_seq[mask_indices]
-            # Get counts only for tokens present in context to avoid giant loops
-            unique_tokens, counts = torch.unique(prev_pred, return_counts=True)
-            for i, tok_id in enumerate(unique_tokens):
-                t_id = tok_id.item()
-                if t_id not in special_ids and counts[i] > 1:
-                    penalty = 5.0 ** (counts[i].item() - 1)
-                    response_logits[:, t_id] = torch.where(
-                        response_logits[:, t_id] < 0,
-                        response_logits[:, t_id] * penalty,
-                        response_logits[:, t_id] / penalty
-                    )
-        
-        if top_k > 0:
-            top_k_vals, _ = torch.topk(response_logits, top_k, dim=-1)
-            min_top_k = top_k_vals[:, -1].unsqueeze(-1)
-            response_logits = response_logits.masked_fill(response_logits < min_top_k, float('-inf'))
+            with torch.no_grad():
+                logits = model(seq_tensor, t, src_key_padding_mask=pad_mask)
+                
+            response_logits = logits[0, mask_indices]
             
-        # Nan protection and sampling
-        if torch.isnan(probs).any():
-            print("WARNING: NaNs detected in probabilities, falling back to uniform")
-            probs = torch.ones_like(probs) / probs.size(-1)
-
-        if step == steps:
-            predicted = torch.argmax(probs, dim=-1)
-        else:
-            try:
-                predicted = torch.multinomial(probs, 1).squeeze(-1)
-            except RuntimeError as e:
-                print(f"Sampling error: {e}")
-                predicted = torch.argmax(probs, dim=-1) # fallback to greedy
+            # Optimized Vectorized Repetition Penalty
+            if step > 1:
+                prev_pred = current_seq[mask_indices]
+                # Get counts only for tokens present in context to avoid giant loops
+                unique_tokens, counts = torch.unique(prev_pred, return_counts=True)
+                for i, tok_id in enumerate(unique_tokens):
+                    t_id = tok_id.item()
+                    if t_id not in special_ids and counts[i] > 1:
+                        penalty = 5.0 ** (counts[i].item() - 1)
+                        response_logits[:, t_id] = torch.where(
+                            response_logits[:, t_id] < 0,
+                            response_logits[:, t_id] * penalty,
+                            response_logits[:, t_id] / penalty
+                        )
             
-        true_probs = F.softmax(response_logits, dim=-1)
-        confidences = true_probs[torch.arange(num_masks), predicted]
-        
-        running_confidence = 0.7 * running_confidence + 0.3 * confidences
-        
-        for i, idx in enumerate(mask_indices):
-            current_seq[idx] = predicted[i]
+            if top_k > 0:
+                top_k_vals, _ = torch.topk(response_logits, top_k, dim=-1)
+                min_top_k = top_k_vals[:, -1].unsqueeze(-1)
+                response_logits = response_logits.masked_fill(response_logits < min_top_k, float('-inf'))
+                
+            probs = F.softmax(response_logits / max(temp, 1e-6), dim=-1)
             
-        if step < steps:
-            target_revealed = int(num_masks * step / steps)
-            num_remask = num_masks - target_revealed
-            if num_remask > 0:
-                _, low_idx = torch.topk(running_confidence, k=num_remask, largest=False)
-                for li in low_idx:
-                    current_seq[mask_indices[li]] = mask_id
-                    
-        seq_tensor = current_seq.unsqueeze(0)
-        out_text = get_output(seq_tensor, is_final=(step == steps))
-        
-        history[-1][1] = out_text
-        yield history
-        time.sleep(0.02)
+            # Nan protection and sampling
+            if torch.isnan(probs).any():
+                print("WARNING: NaNs detected in probabilities, falling back to uniform")
+                probs = torch.ones_like(probs) / probs.size(-1)
+
+            if step == steps:
+                predicted = torch.argmax(probs, dim=-1)
+            else:
+                try:
+                    predicted = torch.multinomial(probs, 1).squeeze(-1)
+                except RuntimeError as e:
+                    print(f"Sampling error: {e}")
+                    predicted = torch.argmax(probs, dim=-1) # fallback to greedy
+                
+            true_probs = F.softmax(response_logits, dim=-1)
+            confidences = true_probs[torch.arange(num_masks), predicted]
+            
+            running_confidence = 0.7 * running_confidence + 0.3 * confidences
+            
+            for i, idx in enumerate(mask_indices):
+                current_seq[idx] = predicted[i]
+                
+            if step < steps:
+                target_revealed = int(num_masks * step / steps)
+                num_remask = num_masks - target_revealed
+                if num_remask > 0:
+                    _, low_idx = torch.topk(running_confidence, k=num_remask, largest=False)
+                    for li in low_idx:
+                        current_seq[mask_indices[li]] = mask_id
+                        
+            seq_tensor = current_seq.unsqueeze(0)
+            out_text = get_output(seq_tensor, is_final=(step == steps))
+            
+            history[-1][1] = out_text
+            yield history
+            time.sleep(0.02)
     except Exception as e:
         import traceback
         error_msg = f"Inference Error: {str(e)}"