build-small-hackathon/InflectionLM

1from transformers import AutoProcessor, AutoModelForCausalLM2import torch3import torch.nn.functional as F4from torch import exp5from typing import Any, Tuple, List6 7def start_model(model_id: str = "google/gemma-4-31B-it"):8    '''9    Initializes and returns the processor and model.10    '''11    processor = AutoProcessor.from_pretrained(model_id)12    model = AutoModelForCausalLM.from_pretrained(13        model_id,14        dtype="auto",15        device_map="auto"16    )17    print(f'Model {model_id} has been installed.')18    return model, processor19 20def make_beams(model: AutoModelForCausalLM, processor: AutoProcessor, initial_prompt: str, temperature: float = 1.0) -> Tuple[Any, List[str]]:21    '''22    Generates 3 diverse responses in response to a prompt.23    '''24    messages = [25        {"role": "system", "content": "You are a helpful assistant."},26        {"role": "user", "content": initial_prompt},27    ]28 29    # Process input30    text = processor.apply_chat_template(31        messages,32        tokenize=False,33        add_generation_prompt=True,34        enable_thinking=False35    )36    inputs = processor(text=text, return_tensors="pt").to(model.device)37 38    # Generate output39    generated_dicts = model.generate(**inputs,40                                    max_new_tokens=1024,41                                    num_beams=1,          # Disable beam search for pure sampling42                                    num_return_sequences=3, # Generate 3 independent diverse samples43                                    return_dict_in_generate=True,44                                    output_scores=True,45                                    temperature=temperature if temperature > 0 else 0.1, # Ensure T > 0 for sampling46                                    top_p=0.9,             # Nucleus sampling for high-quality diversity47                                    top_k=50,              # Top-K sampling to filter noise48                                    do_sample=True)49 50    transcription = processor.batch_decode(generated_dicts.sequences, skip_special_tokens=True)51 52    print('Keys in model output -------------------')53    for key in generated_dicts:54        print(key)55    print('----------------------------------------')56 57    print('Beam scores ----------------------------')58    # sequences_scores is only present in beam search.59    # For sampling, we can approximate the overall score by averaging the transition probabilities.60    if hasattr(generated_dicts, 'sequences_scores'):61        for score in generated_dicts.sequences_scores:62            print(exp(score).item())63    else:64        print('Sampling mode: sequences_scores not available.')65    print('----------------------------------------')66 67    return generated_dicts, transcription68 69def parse_beams(transcription: List[str]) -> List[str]:70    '''71    Parses beams to extract only the response after 'model\nthought'.72    '''73    beam_text = []74    for beam in transcription:75        parts = beam.split('''model\nthought''')76        if len(parts) > 1:77            response = parts[1].strip('\n')78        else:79            response = beam80        beam_text.append(response)81 82    print('Beams have been parsed. --------------')83    return beam_text84 85def get_beam_tokens(generated_dicts: Any, processor: AutoProcessor) -> List[List[str]]:86    '''87    Decodes the generated sequences into individual tokens for each beam.88    '''89    beam_tokens = []90    # The number of generated tokens is the length of the scores list91    gen_len = len(generated_dicts.scores)92 93    for sequence in generated_dicts.sequences:94        # Calculate input length for this specific sequence95        total_len = sequence.shape[0]96        input_len = total_len - gen_len97 98        # Extract only the generated token IDs99        generated_ids = sequence[input_len:].tolist()100 101        # Convert IDs to tokens (e.g., ' Hello', ' world')102        tokens = processor.tokenizer.convert_ids_to_tokens(generated_ids)103        beam_tokens.append(tokens)104 105    return beam_tokens106 107def calculate_score_vectors(model: AutoModelForCausalLM, generated_dicts: Any) -> List[List[float]]:108    '''109    Creates a score vector for each beam containing the probability of each110    token that was chosen.111 112    Optimized to use generated_dicts.scores instead of a full model forward pass113    to prevent GPU timeouts on ZeroGPU.114    '''115    # Number of sequences generated116    num_sequences = generated_dicts.sequences.shape[0]117    # Number of generated tokens (excluding prompt)118    gen_len = len(generated_dicts.scores)119    # Total length of sequences (including prompt)120    total_len = generated_dicts.sequences.shape[1]121    # Input length (prompt length)122    input_len = total_len - gen_len123 124    # Stack the transition scores (logits) from the generation process125    # generated_dicts.scores is a tuple of length gen_len, each element (num_beams, vocab_size)126    all_logits = torch.stack(generated_dicts.scores, dim=0) # shape: (gen_len, num_sequences, vocab_size)127 128    # Convert logits to probabilities across the vocab dimension129    all_probs = F.softmax(all_logits, dim=-1) # shape: (gen_len, num_sequences, vocab_size)130 131    score_vectors = []132    for i in range(num_sequences):133        beam_probs = []134        # Extract the probability for the specific token that was chosen at each step135        for t in range(gen_len):136            token_id = generated_dicts.sequences[i, input_len + t]137            prob = all_probs[t, i, token_id].item()138            beam_probs.append(prob)139        score_vectors.append(beam_probs)140 141    return score_vectors142