natbutter/functiongemma-tuning-lab

1import os2import csv3import json4import shutil5from typing import Optional, List, Any6from huggingface_hub import login7from transformers import AutoTokenizer, AutoModelForCausalLM8from tools import DEFAULT_SYSTEM_MSG 9# Note: We do NOT import TOOLS here anymore to avoid stale data10 11def authenticate_hf(token: Optional[str]) -> None:12    """Logs into the Hugging Face Hub."""13    if token:14        print("Logging into Hugging Face Hub...")15        login(token=token)16    else:17        print("Skipping Hugging Face login: HF_TOKEN not set.")18 19def load_model_and_tokenizer(model_name: str):20    print(f"Loading Transformer model: {model_name}")21    try:22        target_model = model_name23        if model_name.startswith("..") and not os.path.exists(model_name):24            print(f"Warning: Local path {model_name} not found. Falling back to default hub model.")25            target_model = "google/gemma-2b-it" 26 27        tokenizer = AutoTokenizer.from_pretrained(target_model)28        model = AutoModelForCausalLM.from_pretrained(target_model)29        print("Model loaded successfully.")30        return model, tokenizer31    except Exception as e:32        print(f"Error loading Transformer model {target_model}: {e}")33        raise e34 35# UPDATED: Now accepts tools_list as an argument36def create_conversation_format(sample, tools_list):37    """Formats a dataset row into the conversational format required for SFT."""38    try:39        tool_args = json.loads(sample["tool_arguments"])40    except (json.JSONDecodeError, TypeError):41        tool_args = {}42        43    return {44        "messages": [45            {"role": "developer", "content": DEFAULT_SYSTEM_MSG},46            {"role": "user", "content": sample["user_content"]},47            {"role": "assistant", "tool_calls": [{"type": "function", "function": {"name": sample["tool_name"], "arguments": tool_args}}]},48        ],49        "tools": tools_list # Injects the dynamic tools50    }51 52def parse_csv_dataset(file_path: str) -> List[List[str]]:53    """Parses an uploaded CSV file."""54    dataset = []55    if not file_path:56        return dataset57        58    with open(file_path, 'r', newline='', encoding='utf-8') as f:59        reader = csv.reader(f)60        try:61            header = next(reader)62            if not (header and "user_content" in header[0].lower()):63                f.seek(0)64        except StopIteration:65            return dataset66 67        for row in reader:68            if len(row) >= 3:69                dataset.append([s.strip() for s in row[:3]])70    return dataset71 72def zip_directory(source_dir: str, output_name_base: str) -> str:73    """Zips a directory."""74    return shutil.make_archive(75        base_name=output_name_base,76        format='zip',77        root_dir=source_dir,78    )79