dassarthak18/phi2-fstar-lora

Phi-2 Fine-tuned on F* (LoRA)

This model is a LoRA fine-tuned version of Microsoft's Phi-2 trained to generate F* code and automated proofs in the F* formal verification language. In their paper Towards Neural Synthesis forSMT-Assisted Proof-Oriented Programming, Saikat Chakraborty et al. (2024) use Phi-2, among other LLMs, for evaluating their dataset. This LLM is relatively small and can be fine-tuned on consumer-grade GPUs using LoRA.

Model Details

Model Description

This is a parameter-efficient fine-tuned model using LoRA (Low-Rank Adaptation) on the Phi-2 base model. It specializes in generating F* code, including function definitions, lemmas, and formal proofs.

• Developed by: Sarthak Das
• Model type: Causal Language Model (LoRA Adapter)
• Language(s): F* (Formal verification language)
• License: MIT License
• Finetuned from model: microsoft/phi-2
• Training dataset: microsoft/FStarDataSet-V2

Model Sources

• Base Model Repository: microsoft/phi-2
• Dataset: microsoft/FStarDataSet-V2

Training Setup

• Optimizer: AdamW
• Adapter: LoRA (r=8, α=16, dropout=0.05)
• Quantization: None (full fp16 model)
• Precision: fp16 mixed precision (trained on RTX 5060 8GB VRAM)
• Batch size: 1 per device (effective batch size: 16 with gradient accumulation)
• Epochs: 2
• LR: 2e-4

The training script is available in the ./src directory for reproduction, and can be run via:

cd src
chmod u+x train.sh
./train.sh

Evaluation

To examine the tensorboard logs for further training and eval details, use the following commands:

pip install tensorboard
tensorboard --logdir ./logs

Testing Data & Metrics

• Validation Loss: 0.5417
• Validation Perplexity: 1.72

How to Get Started with the Model

Installation

pip install transformers peft torch

Basic Usage

from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel
import torch

# Load base model
base_model = AutoModelForCausalLM.from_pretrained(
    "microsoft/phi-2",
    torch_dtype=torch.float16,
    device_map="auto",
    trust_remote_code=True
)

# Load LoRA adapter
model = PeftModel.from_pretrained(base_model, "dassarthak18/phi2-fstar-lora")
tokenizer = AutoTokenizer.from_pretrained("microsoft/phi-2", trust_remote_code=True)

# Generate F* code
prompt = "let rec factorial (n: nat) : nat ="
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)

outputs = model.generate(
    **inputs,
    max_new_tokens=512,
    do_sample=False,
    num_beams=3,
    pad_token_id=tokenizer.eos_token_id,
    early_stopping=True,
)

generated_code = tokenizer.decode(outputs[0], skip_special_tokens=True)
if "\\n\\n" in generated_code:
    generated_code = generated_code.split("\\n\\n")[0]
elif "\n\n" in generated_code:
    generated_code = generated_code.split("\n\n")[0]
print(generated_code)

Example Outputs

Prompt:

let rec factorial (n: nat) : nat =

Generated:

let rec factorial (n: nat) : nat =
    if n = 0 then 1 else n * factorial (n-1)

Prompt:

module Example
open FStar.Mul

Generated:

module Example
open FStar.Mul

let lemma_multiplication (a: nat) (b: nat) : 
    Lemma (a * b = b * a) = ()