llmflowstack 1.0.0__py3-none-any.whl

llmflowstack/models/LLaMA3.py

@@ -0,0 +1,213 @@
+import textwrap
+import threading
+from time import time
+from typing import Iterator, Literal, TypedDict, cast
+
+import torch
+from transformers import (AutoTokenizer, StoppingCriteriaList,
+                          TextIteratorStreamer)
+from transformers.models.llama import LlamaForCausalLM
+from transformers.utils.quantization_config import BitsAndBytesConfig
+
+from llmflowstack.base.base import BaseModel
+from llmflowstack.callbacks.stop_on_token import StopOnToken
+from llmflowstack.schemas.params import GenerationParams
+from llmflowstack.utils.exceptions import MissingEssentialProp
+from llmflowstack.utils.generation_utils import create_generation_params
+
+
+class LLaMA3Input(TypedDict):
+	input_text: str
+	expected_answer: str | None
+	system_message: str | None
+
+class LLaMA3(BaseModel):
+	model: LlamaForCausalLM | None = None
+	question_fields = ["input_text", "system_message"]
+	answer_fields = ["expected_answer"]
+
+	def _set_generation_stopping_tokens(
+		self,
+		tokens: list[int]
+	) -> None:
+		if not self.tokenizer:
+			self._log("Could not set stop tokens - generation may not work...", "WARNING")
+			return None
+		particular_tokens = self.tokenizer.encode("<|eot_id|>")
+		self.stop_token_ids = tokens + particular_tokens
+
+	def _load_model(
+		self,
+		checkpoint: str,
+		quantization: Literal["8bit", "4bit"] | bool | None = None
+	) -> None:
+		quantization_config = None
+		if quantization == "4bit":
+			quantization_config = BitsAndBytesConfig(
+				load_in_4bit=True
+			)
+			self.model_is_quantized = True
+		if quantization == "8bit":
+			quantization_config = BitsAndBytesConfig(
+				load_in_8bit=True
+			)
+			self.model_is_quantized = True
+
+		self.model = LlamaForCausalLM.from_pretrained(
+			checkpoint,
+			quantization_config=quantization_config,
+			dtype="auto",
+			device_map="auto",
+			attn_implementation="eager"
+		)
+
+	def _build_input(
+		self,
+		input_text: str,
+		expected_answer: str | None = None,
+		system_message: str | None = None
+	) -> str:
+		if not self.tokenizer:
+			raise MissingEssentialProp("Could not find tokenizer.")
+
+		answer = f"{expected_answer}{self.tokenizer.eos_token}" if expected_answer else ""
+
+		return textwrap.dedent(
+			f"<|start_header_id|>system<|end_header_id|>{system_message or ""}\n"
+			f"<|eot_id|><|start_header_id|>user<|end_header_id|>{input_text}\n"
+			f"<|eot_id|><|start_header_id|>assistant<|end_header_id|>{answer}"
+		)
+
+	def build_input(
+		self,
+		input_text: str,
+		system_message: str | None = None,
+		expected_answer: str | None = None
+	) -> LLaMA3Input:
+		if not self.tokenizer:
+			raise MissingEssentialProp("Could not find tokenizer.")
+
+		return {
+			"input_text": input_text,
+			"system_message": system_message,
+			"expected_answer": expected_answer
+		}
+
+	def generate(
+		self,
+		input: LLaMA3Input | str,
+		params: GenerationParams | None = None
+	) -> str | None:
+		if self.model is None or self.tokenizer is None:
+			self._log("Model or Tokenizer missing", "WARNING")
+			return None
+
+		self.model
+
+		self._log(f"Processing received input...'")
+
+		if params is None:
+			params = GenerationParams(max_new_tokens=8192)
+		elif params.max_new_tokens is None:
+			params.max_new_tokens = 8192
+
+		generation_params = create_generation_params(params)
+		self.model.generation_config = generation_params
+
+		if params:
+			generation_params = create_generation_params(params)
+			self.model.generation_config = generation_params
+
+		model_input = None
+		if isinstance(input, str):
+			model_input = self._build_input(
+				input_text=input
+			)
+		else:
+			model_input = self._build_input(
+				input_text=input["input_text"],
+				system_message=input.get("system_message", "")
+			)
+
+		tokenized_input = self._tokenize(model_input)
+
+		input_ids, attention_mask = tokenized_input
+
+		self.model.eval()
+		self.model.gradient_checkpointing_disable()
+
+		start = time()
+
+		with torch.no_grad():
+			outputs = self.model.generate(
+				input_ids=input_ids,
+				attention_mask=attention_mask,
+				use_cache=True,
+				eos_token_id=None,
+				stopping_criteria=StoppingCriteriaList([StopOnToken(self.stop_token_ids)])
+			)
+
+		end = time()
+		total_time = end - start
+
+		self._log(f"Response generated in {total_time:.4f} seconds")
+
+		response = outputs[0][input_ids.shape[1]:]
+
+		return self.tokenizer.decode(response, skip_special_tokens=True)
+	
+	def generate_stream(
+		self,
+		input: LLaMA3Input | str,
+		params: GenerationParams | None = None
+	) -> Iterator[str]:
+		if self.model is None or self.tokenizer is None:
+			self._log("Model or Tokenizer missing", "WARNING")
+			if False:
+				yield ""
+			return
+		
+		if params is None:
+			params = GenerationParams(max_new_tokens=8192)
+		elif params.max_new_tokens is None:
+			params.max_new_tokens = 8192
+
+		generation_params = create_generation_params(params)
+		self.model.generation_config = generation_params
+
+		if isinstance(input, str):
+			model_input = self._build_input(
+				input_text=input
+			)
+		else:
+			model_input = self._build_input(
+				input_text=input["input_text"],
+				system_message=input.get("system_message")
+			)
+		
+		tokenized_input = self._tokenize(model_input)
+		input_ids, attention_mask = tokenized_input
+
+		streamer = TextIteratorStreamer(
+			cast(AutoTokenizer, self.tokenizer),
+			skip_prompt=True,
+			skip_special_tokens=True
+		)
+
+		def _generate() -> None:
+			assert self.model is not None
+			with torch.no_grad():
+				self.model.generate(
+					input_ids=input_ids,
+					attention_mask=attention_mask,
+					use_cache=True,
+					eos_token_id=None,
+					streamer=streamer,
+					stopping_criteria=StoppingCriteriaList([StopOnToken(self.stop_token_ids)])
+				)
+		
+		thread = threading.Thread(target=_generate)
+		thread.start()
+
+		for new_text in streamer:
+			yield new_text

llmflowstack/models/__init__.py

@@ -0,0 +1,9 @@
+from .Gemma import Gemma
+from .GPT_OSS import GPT_OSS
+from .LLaMA3 import LLaMA3
+
+__all__ = [
+  "Gemma",
+  "GPT_OSS",
+  "LLaMA3"
+]

llmflowstack/rag/__iinit__.py

@@ -0,0 +1,5 @@
+from .pipeline import RAGPipeline
+
+__all__ = [
+  "RAGPipeline"
+]

llmflowstack/rag/pipeline.py

@@ -0,0 +1,114 @@
+import uuid
+
+from langchain_chroma import Chroma
+from langchain_core.documents import Document
+from langchain_core.embeddings import Embeddings
+from langchain_text_splitters import RecursiveCharacterTextSplitter
+from sentence_transformers import SentenceTransformer
+
+
+class EncoderWrapper(Embeddings):
+	def __init__(
+		self,
+		model: SentenceTransformer
+	) -> None:
+		self.model = model
+
+	def embed_documents(
+		self,
+		texts: list[str]
+	) -> list[list[float]]:
+		return self.model.encode(texts, task="retrieval", show_progress_bar=True).tolist()
+	
+	def embed_query(
+		self,
+		text: str
+	) -> list[float]:
+		return self.model.encode(text, task="retrieval", show_progress_bar=True).tolist()
+
+class RAGPipeline:
+	def __init__(
+		self,
+		checkpoint: str,
+		collection_name: str = "rag_memory",
+		persist_directory: str = "./chroma_store",
+		chunk_size: int = 1000,
+		chunk_overlap: int = 200
+	) -> None:
+		
+		self.encoder = SentenceTransformer(checkpoint, trust_remote_code=True)
+
+		self.vector_store = Chroma(
+			collection_name=collection_name,
+			embedding_function=EncoderWrapper(self.encoder),
+			persist_directory=persist_directory
+		)
+
+		self.splitter = RecursiveCharacterTextSplitter(
+			chunk_size=chunk_size,
+			chunk_overlap=chunk_overlap,
+			add_start_index=True,
+		)
+	
+	def index_documents(
+		self,
+		docs: list[Document],
+		ids: list[str]
+	) -> None:
+		splits = self.splitter.split_documents(docs)
+		split_ids = [f"{ids[0]}_{i}" for i in range(len(splits))]
+		self.vector_store.add_documents(splits, ids=split_ids)
+	
+	def create(
+		self,
+		information: str,
+		other_info: dict[str, str] = {},
+		doc_id: str | None = None,
+		should_index: bool = True
+	) -> Document:
+		if doc_id is None:
+			doc_id = str(uuid.uuid4())
+			
+		doc = Document(
+			page_content=information,
+			metadata={"id": doc_id, **other_info}
+		)
+
+		if should_index:
+			self.index_documents([doc], ids=[doc_id])
+
+		return doc
+	
+	def update(
+		self,
+		doc_id: str,
+		new_information: str,
+		other_info: dict[str, str] = {}
+	) -> Document:
+		self.vector_store.delete(ids=[doc_id])
+
+		return self.create(
+			information=new_information,
+			other_info=other_info,
+			doc_id=doc_id
+		)
+	
+	def delete(
+		self, doc_id: str
+	) -> None:
+		self.vector_store.delete(ids=[doc_id])
+
+	def query(
+		self,
+		query: str,
+		k: int = 4,
+		category: str | None = None
+	) -> str:
+		if category:
+			docs = self.vector_store.similarity_search(
+				query, k=k, filter={"category": category}
+			)
+		else:
+			docs = self.vector_store.similarity_search(query, k=k)
+
+		return "\n\n".join(doc.page_content for doc in docs)

llmflowstack/schemas/__init__.py

@@ -0,0 +1,9 @@
+from .params import (GenerationBeamsParams, GenerationParams,
+                     GenerationSampleParams, TrainParams)
+
+__all__ = [
+  "GenerationBeamsParams",
+  "GenerationParams",
+  "GenerationSampleParams",
+  "TrainParams"
+]

llmflowstack/schemas/params.py

@@ -0,0 +1,39 @@
+from dataclasses import dataclass, field
+from typing import Literal
+
+from transformers import TextIteratorStreamer
+
+
+@dataclass
+class TrainParams:
+  batch_size: int = 1
+  gradient_accumulation: int = 8
+  epochs: int = 1
+  warmup_ratio: float = 0.0
+  lr: float = 2e-5
+  optim: Literal[
+    "adamw_torch",
+    "adamw_torch_fused",
+    "sgd"
+  ] = "adamw_torch"
+  logging_steps=1
+
+@dataclass
+class GenerationBeamsParams:
+  num_beams: int | None = None
+  length_penalty: float | None = None
+  early_stopping: bool | None = None
+
+@dataclass
+class GenerationSampleParams:
+  temperature: float | None = None
+  top_p: float | None = None
+  typical_p: float | None = None
+
+@dataclass
+class GenerationParams:
+  max_new_tokens: int | None = None
+  repetition_penalty: float | None = None
+  sample: GenerationSampleParams = field(default_factory=GenerationSampleParams)
+  beams: GenerationBeamsParams = field(default_factory=GenerationBeamsParams)
+  streamer: TextIteratorStreamer | None = None

llmflowstack/utils/__init__.py

@@ -0,0 +1,11 @@
+from .evaluation_methods import (bert_score_evaluation,
+                                 cosine_similarity_evaluation,
+                                 rouge_evaluation, text_evaluation)
+
+__all__ = [
+  "bert_score_evaluation",
+  "cosine_similarity_evaluation",
+  "rouge_evaluation",
+  "evaluation_methods",
+  "text_evaluation"
+]

llmflowstack/utils/evaluation_methods.py

@@ -0,0 +1,92 @@
+from evaluate import load
+from nltk.stem.snowball import SnowballStemmer
+from rouge_score import rouge_scorer
+from sentence_transformers import SentenceTransformer, util
+
+
+def stem_texts(texts: list[str]) -> list[str]:
+	stemmer = SnowballStemmer("portuguese")
+
+	stemmed_texts: list[str] = []
+	for text in texts:
+		stemmed_text = " ".join([stemmer.stem(word) for word in text.split()])
+		stemmed_texts.append(stemmed_text)
+
+	return stemmed_texts
+
+def rouge_evaluation(
+	preds: list[str],
+	refs: list[str]
+) -> dict[str, float]:
+	preds_stemmed = stem_texts(preds)
+	refs_stemmed = stem_texts(refs)
+
+	rouge_metrics = {"rouge1": [], "rouge2": [], "rougeL": []}
+	scorer = rouge_scorer.RougeScorer(["rouge1", "rouge2", "rougeL"], use_stemmer=False)
+
+	for ref, pred in zip(refs_stemmed, preds_stemmed):
+		scores = scorer.score(ref, pred)
+		for key in rouge_metrics:
+			rouge_metrics[key].append(scores[key].fmeasure)
+
+	return {k: sum(v)/len(v) for k, v in rouge_metrics.items()}
+
+def bert_score_evaluation(
+	preds: list[str],
+	refs: list[str]
+) -> dict[str, float]:
+	bertscore = load("bertscore")
+
+	bert_result = bertscore.compute(predictions=preds, references=refs, lang="pt")
+
+	bert_avg = {}
+	if bert_result:
+		bert_avg = {
+			"bertscore_precision": sum(bert_result["precision"]) / len(bert_result["precision"]),
+			"bertscore_recall": sum(bert_result["recall"]) / len(bert_result["recall"]),
+			"bertscore_f1": sum(bert_result["f1"]) / len(bert_result["f1"])
+		}
+
+	return bert_avg
+
+def cosine_similarity_evaluation(
+	preds: list[str],
+	refs: list[str]
+) -> dict[str, float]:
+	model = SentenceTransformer("sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2")
+
+	emb_preds = model.encode(preds, convert_to_tensor=True)
+	emb_refs = model.encode(refs, convert_to_tensor=True)
+
+	cos_sim_matrix = util.cos_sim(emb_preds, emb_refs)
+
+	cos_sim_scores = cos_sim_matrix.diag()  
+	avg_cos_sim = cos_sim_scores.mean().item()
+
+	return {"cosine_similarity": float(avg_cos_sim)}
+
+def text_evaluation(
+	preds: list[str],
+	refs: list[str],
+	rouge: bool = True,
+	bert: bool = True,
+	cosine: bool = True
+) -> dict[str, float]:
+	result = {}
+	if rouge:
+		result.update(rouge_evaluation(
+			preds=preds,
+			refs=refs
+		))
+	if bert:
+		result.update(bert_score_evaluation(
+			preds=preds,
+			refs=refs
+		))
+	if cosine:
+		result.update(cosine_similarity_evaluation(
+			preds=preds,
+			refs=refs
+		))
+
+	return result

llmflowstack/utils/exceptions.py

@@ -0,0 +1,2 @@
+class MissingEssentialProp(Exception):
+  pass

llmflowstack/utils/generation_utils.py

@@ -0,0 +1,30 @@
+from transformers.generation.configuration_utils import GenerationConfig
+
+from llmflowstack.schemas.params import GenerationParams
+
+
+def create_generation_params(generation_configs: GenerationParams) -> GenerationConfig:
+
+	params = {
+		"max_new_tokens": generation_configs.max_new_tokens,
+		"repetition_penalty": generation_configs.repetition_penalty
+	}
+	if generation_configs.sample:
+		sample = generation_configs.sample
+		params.update({
+			"do_sample": True,
+			"temperature": sample.temperature,
+			"top_p": sample.top_p,
+			"typical_p": sample.typical_p,
+			"num_beams": 1
+		})
+	elif generation_configs.beams == "beams":
+		beams = generation_configs.beams
+		params.update({
+			"do_sample": False,
+			"num_beams": beams.num_beams,
+			"length_penalty": beams.length_penalty,
+			"early_stopping": beams.early_stopping
+		})
+
+	return GenerationConfig(**params)