Vaan-V1_400M

Vaan-V1_400M

Developed by Suraj Kumar Mallick, Vaan-V1_400M is a causal language model built on the classic GPT-2 architecture. This model was designed to be a lightweight yet capable assistant, balancing efficiency with a solid knowledge base.

Model Details

• Model Type: GPT-2 (Causal Language Modeling)
• Parameters: ~400 Million
• Architecture: Standard GPT-2 decoder-only transformer
• Developer: Suraj Kumar Mallick
• Language(s): English
• License: Apache 2.0

Training Information

The training of Vaan-V1 followed a multi-stage process to ensure broad general knowledge followed by task-specific instruction following.

1. Pre-training

The model was pre-trained on a massive corpus to develop linguistic intuition and factual knowledge:

• Wikipedia Dumps: For structured, encyclopedic knowledge.
• RedPajama: Used to diversify the model's understanding across web data, books, and code.

2. Fine-tuning

The model underwent Instruction Fine-Tuning (IFT) using:

• Dataset: tatsu-lab/alpaca
• Goal: To enable the model to follow user prompts and handle conversational tasks effectively.

Hardware

• GPU: Trained on a single NVIDIA T4 GPU.
• Training Framework: Hugging Face transformers and accelerate.

Intended Use

Vaan-V1_400M is suitable for:

• Text generation and creative writing.
• Basic instruction following.
• Educational purposes regarding small-scale LLM behavior.
• Deployment on edge devices with limited VRAM.

How to Use

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer
from threading import Thread

model_id = "surajpuku/Vaan-V1_400M"

# Load tokenizer + model
tokenizer = AutoTokenizer.from_pretrained(model_id)

model = AutoModelForCausalLM.from_pretrained(
    model_id,
    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
    device_map="auto" if torch.cuda.is_available() else None
)

if not torch.cuda.is_available():
    model.to("cpu")

def stream_generate(prompt):
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)

    streamer = TextIteratorStreamer(
        tokenizer,
        skip_prompt=True,
        skip_special_tokens=True
    )

    generation_kwargs = dict(
        **inputs,
        max_new_tokens=256,
        temperature=0.7,
        top_p=0.9,
        do_sample=True,
        streamer=streamer
    )

    thread = Thread(target=model.generate, kwargs=generation_kwargs)
    thread.start()

    for token in streamer:
        print(token, end="", flush=True)

    print()


print("Chat started. Type 'exit' to quit.")

while True:
    user_input = input("\nYou: ")
    if user_input.lower() in ["exit", "quit"]:
        break

    print("Assistant: ", end="", flush=True)
    stream_generate(user_input)