Lambent/

Qwen3-4B-Base-Continued-GRPO

class RewardModel:
    """Multi-domain reward model for GRPO pretraining."""

    def __init__(self, device: str = "cuda"):
        self.device = device

        # Load components
        print("Loading reward model components...")
        self.semantic_model = SentenceTransformer('all-MiniLM-L6-v2', device=device)
        self.chrf = CHRF(word_order=2)
        self.rouge = rouge_scorer.RougeScorer(['rougeL'], use_stemmer=True)

    def _levenshtein_distance(self, s1: str, s2: str) -> int:
        """C-optimized Levenshtein distance."""
        return Levenshtein.distance(s1, s2)

    def compute_reward(
        self,
        prediction: str,
        reference: str,
        reward_type: str
    ) -> float:
        """Compute reward for a single prediction."""

        if not prediction or not reference:
            return 0.0

        if reward_type == "creative":
            # Semantic + chrF++
            embs = self.semantic_model.encode([reference, prediction], convert_to_tensor=True)
            sem_score = torch.nn.functional.cosine_similarity(embs[0:1], embs[1:2]).item()
            chrf_score = self.chrf.sentence_score(prediction, [reference]).score / 100.0
            return 0.4 * sem_score + 0.6 * chrf_score

        elif reward_type == "hybrid":
            # Semantic + ROUGE-L + length penalty
            embs = self.semantic_model.encode([reference, prediction], convert_to_tensor=True)
            sem_score = torch.nn.functional.cosine_similarity(embs[0:1], embs[1:2]).item()
            rouge_result = self.rouge.score(reference, prediction)
            rouge_l = rouge_result['rougeL'].fmeasure

            ref_len = len(reference.split())
            pred_len = len(prediction.split())
            if ref_len > 0:
                len_ratio = pred_len / ref_len
                length_penalty = max(0.0, 1.0 - abs(1.0 - len_ratio))
            else:
                length_penalty = 1.0 if pred_len == 0 else 0.0

            return 0.6 * sem_score + 0.3 * rouge_l + 0.1 * length_penalty

        elif reward_type == "levenshtein":
            # Edit distance
            max_len = max(len(prediction), len(reference))
            if max_len == 0:
                return 1.0
            dist = self._levenshtein_distance(prediction, reference)
            return max(0.0, 1.0 - (dist / max_len))

        else:
            raise ValueError(f"Unknown reward type: {reward_type}")