llmflowstack 1.2.6
LLMFlowStack is a framework for training and using LLMs (LLaMA, GPT-OSS, Gemma, ...). Supports DAPT, fine-tuning, and distributed inference. Public fork without institution-specific components.
License: MIT
Requires Python: >=3.12

# LLMFlowStack

**LLMFlowStack** is a lightweight framework designed to simplify the use of LLMs (LLaMA, GPT-OSS, and Gemma) for NLP tasks.

> **Note:** LLMFlowStack is intended for high-performance machines with **one or more NVIDIA H100 GPUs**.

It provides:

- **Training pipelines** with **fine-tuning** or **DAPT** in distributed setups — just provide the data and the process runs automatically;
- **Distributed inference** made simple;
- **Evaluation** with standard metrics (BERTScore, ROUGE, Cosine Similarity).

The goal is to make experimentation with LLMs more accessible, without the need to build complex infrastructure from scratch.

## Supported Models

This framework is designed to provide flexibility when working with different open-source and commercial LLMs. Currently, the following models are supported:

- **GPT-OSS**

  - [`GPT-OSS 20B`](https://huggingface.co/openai/gpt-oss-20b)
  - [`GPT-OSS 120B`](https://huggingface.co/openai/gpt-oss-120b)
    > Fine-Tuning, DAPT and Inference Available

- **LLaMA 3**

  - [`LLaMA 3.1 8B - Instruct`](https://huggingface.co/meta-llama/Llama-3.1-8B-Instruct)
  - [`LLaMA 3.1 70B - Instruct`](https://huggingface.co/meta-llama/Llama-3.1-70B-Instruct)
  - [`LLaMA 3.3 70B - Instruct`](https://huggingface.co/meta-llama/Llama-3.3-70B-Instruct)
  - [`LLaMA 3.3 405B - Instruct`](https://huggingface.co/meta-llama/Llama-3.1-405B-Instruct)
    > Fine-Tuning, DAPT and Inference Available

- **LLaMA 4**

  - [`LLaMA 4 Scout - Instruct`](https://huggingface.co/meta-llama/Llama-4-Scout-17B-16E-Instruct)
    > DAPT and Inference Available

- **Gemma**

  - [`Gemma 3 27B - Instruct`](https://huggingface.co/google/gemma-3-27b-it)
    > DAPT and Inference Available

- **MedGemma**
  - [`MedGemma 27B Text - Instruct`](https://huggingface.co/google/medgemma-27b-text-it)
    > Fine-Tuning, DAPT and Inference Available

> Other architectures based on those **may** function correctly.

---

## Installation

You can install the package directly from [PyPI](https://pypi.org/project/llmflowstack/):

```bash
pip install llmflowstack
```

## Usage

This section presents a bit of what you can do with the framework.

### Loading models

You can load as many models as your hardware allows (H100 GPU recommended)...

```python
from llmflowstack import GPT_OSS, LLaMA3

# Loading a LLaMA model
first_model = LLaMA3()
first_model.load_checkpoint(
  checkpoint="/llama-3.1-8b-Instruct",
)

# Loading a quantized LLaMA model
second_model = LLaMA3(
  checkpoint="/llama-3.3-70b-Instruct",
  quantization="4bit"
)

# Loading a GPT-OSS, quantized and with seed
thrid_model = GPT_OSS(
  checkpoint="/gpt-oss-20b",
  quantization=True,
  seed=1234
)
```

### Inference Examples

```python
> from llmflowstack import GPT_OSS, GenerationParams, GenerationSampleParams

> gpt_oss_model = GPT_OSS(checkpoint="/gpt-oss-120b")

> gpt_oss_model.generate("Tell me a joke!")
'Why did the scarecrow become a successful motivational speaker? Because he was outstanding **in** his field! 🌾😄'

# Exclusive for GPT-OSS
> gpt_oss_model.set_reasoning_level("High")

> custom_input = gpt_oss_model.build_input(
    input_text="Tell me another joke!",
    developer_message="You are a clown and after every joke, you should say 'HONK HONK'"
  )
> gpt_oss_model.generate(
    input=custom_input,
    params=GenerationParams(
      max_new_tokens=1024,
      sample=GenerationSampleParams(
        temperature=0.3
      )
    )
  )
'Why did the scarecrow win an award? Because he was outstanding in his field!  \n\nHONK HONK'

> llama_model = LLaMA3(checkpoint="/llama-3.3-70B-Instruct", quantization="4bit")
> llama_model.generate("Why is the sky blue?")
'The sky appears blue because of a phenomenon called Rayleigh scattering, which is the scattering of light'

# You can also disable GPT-OSS reasoning, but this works only when the model is being used strictly for inference. If the model has been trained or fine-tuned beforehand, this option will not behave correctly.
> gpt_oss_model.set_reasoning_level("Off") # (inference-only)
```

You can also generate tokens using a streamer, that is, receiving one token at a time by using the iterator version of the generate function:

```python
llama_4 = LLaMA4(
  checkpoint="llama-4-scout-17b-16e-instruct"
)

it = llama_4.generate_stream("Who was Alan Turing?")

for text in it:
  print(text, end="", sep="")   # The model will keep yielding tokens until it reaches an end-of-generation token (or until you stop iterating)
```

### Training Examples (DAPT & Fine-tune)

```python
from llmflowstack import LLaMA3
from llmflowstack.schemas import TrainParams

model = LLaMA3(
  checkpoint="llama-3.1-8b-Instruct"
)

# Creating the dataset
dataset = []
dataset.append(model.build_input(
  input_text="Chico is a cat, which color he is?",
  expected_answer="Black!"
))

dataset.append(model.build_input(
  input_text="Fred is a dog, which color he is?",
  expected_answer="White!"
))

# Does the DAPT in the full model
model.dapt(
  train_dataset=dataset,
  params=TrainParams(
    batch_size=1,
    epochs=3,
    gradient_accumulation=1,
    lr=2e-5
  )
)

# Does the fine-tune this time
model.fine_tune(
  train_dataset=dataset,
  params=TrainParams(
    batch_size=1,
    gradient_accumulation=1,
    lr=2e-5,
    epochs=50
  ),
  save_at_end=True,
  # It will save the model
  save_path="./output"
)

# Saving the final result
model.save_checkpoint(
  path="./model-output"
)
```

### RAG Pipeline

A prototype of a RAG pipeline is also available. You can instantiate and use it as follows:

```python
from llmflowstack import VectorDatabase

vector_db = VectorDatabase(
	checkpoint="jina-embeddings-v4",
	chunk_size=1000,
	chunk_overlap=200
)

# Create or load an existing collection
vector_db.get_collection(
	collection_name="memory_rag",
	persist_directory="./memory"
)

vector_db.get_collection(
	collection_name="files_rag",
	persist_directory="./files"
)

# You may also omit the persist directory; in this case, the RAG data will be stored in memory
vector_db.get_collection(
	collection_name="files_rag"
)

# To create a new document in a collection
vector_db.create(
	collection_name="memory_rag",
	information="User loves Pizza!",    # Main information to be indexed in the vector database
	other_info={"category": "food"},
	can_split=False,                    # Indicates whether the information can be split into chunks
	should_index=True                   # Defaults to True — defines whether the document should be indexed or only returned as a Document instance
)

# After adding documents, you can query the database
query_result = vector_db.query(
	collection_name="memory_rag",
	query="pizza",
	filter={"category": "food"},
	k=3   # Number of chunks to retrieve
)

print(query_result)
# > "User loves Pizza!"
```

### NLP Evaluation

> **Disclaimer**
> These evaluation functions are designed for batch processing. Models and encoders are loaded internally on each call, which may be inefficient for per-sample or streaming evaluation.

```python
> from llmflowstack import text_evaluation
> from llmflowstack.utils import (bert_score_evaluation, bleu_score_evaluation, cosine_similarity_evaluation, rouge_evaluation)

# Predictions from some model
> predictions = ["Chico is a dog, and he is orange!", "Fred is a cat, and he is white!"]
# References text (ground truth)
> references = ["Chico is a cat, and he is black!", "Fred is a dog, and he is white!"]

# BERT Score Evaluation
> bert_score_evaluation(predictions, references)
{'bertscore_precision': 0.9773, 'bertscore_recall': 0.9773, 'bertscore_f1': 0.9773}

# Bleu Score Evaluation
> bleu_score_evaluation(predictions, references)
{'bleu_score': 0.3656}

# Cosine Similarity Evaluation
> cosine_similarity_evaluation(predictions, references)
{'cosine_similarity': 0.7443}

# Rouge Score Evaluation
> rouge_evaluation(predictions, references)
{'rouge1': 0.8125, 'rouge2': 0.6429, 'rougeL': 0.8125}

# All-in-one function
> text_evaluation(predictions, references)
{'bertscore_precision': 0.9773, 'bertscore_recall': 0.9773, 'bertscore_f1': 0.9773, 'bleu_score': 0.3656, 'cosine_similarity': 0.7443, 'rouge1': 0.8125, 'rouge2': 0.6429, 'rougeL': 0.8125}
```