Document question answering with Llama 3 and Weaviate

!pip install torch transformers weaviate-client langchain datasets bitsandbytes wandb

import wandb
wandb.login()

from datasets import load_dataset
dataset = load_dataset("squad", split="train")

from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
import torch

# Load the tokenizer and model with QLoRA configuration
compute_dtype = torch.bfloat16  # Set to bfloat16 if supported, else use float16
bnb_config = BitsAndBytesConfig(load_in_4bit=True)

# Load Llama 3.1 405B with quantization config
model_name = "NousResearch/Llama-3-405B"
model = AutoModelForCausalLM.from_pretrained(model_name, quantization_config=bnb_config)

# Load Llama tokenizer
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
tokenizer.pad_token = tokenizer.eos_token  # Avoid overflow issue

# Optional: Fine-tune LoRA configuration for QA task
from peft import LoraConfig, SFTTrainer, TrainingArguments

peft_config = LoraConfig(r=8, lora_alpha=32, lora_dropout=0.1)
training_args = TrainingArguments(
    output_dir="./outputs",
    num_train_epochs=3,
    per_device_train_batch_size=8,
    gradient_accumulation_steps=4,
    save_steps=100,
    logging_steps=10,
    learning_rate=2e-5,
    bf16=True,  # If GPU supports bf16, otherwise use fp16
    report_to="wandb"  # Log to W&B
)

# Set up trainer for supervised fine-tuning
trainer = SFTTrainer(
    model=model,
    train_dataset=dataset,
    peft_config=peft_config,
    args=training_args,
    tokenizer=tokenizer
)

trainer.train()

# Save trained model
trainer.model.save_pretrained(new_model)

import weaviate

client = weaviate.Client(embedded_options=weaviate.EmbeddedOptions())

# Load your documents and split into manageable chunks
from langchain.document_loaders import PyPDFLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter

loader = PyPDFLoader("path_to_your_document.pdf")
docs = loader.load()

text_splitter = RecursiveCharacterTextSplitter(chunk_size=1500, chunk_overlap=100)
splits = text_splitter.split_documents(docs)

# Store chunks in Weaviate
from langchain.vectorstores import Weaviate
from langchain.embeddings import HuggingFaceEmbeddings

vectorstore = Weaviate.from_documents(
    client=client,
    documents=splits,
    embedding=HuggingFaceEmbeddings(model_name="NousResearch/Llama-3-405B")
)

dataset = load_dataset(dataset_name, split="train")

# Load tokenizer and model with QLoRA configuration
compute_dtype = getattr(torch, bnb_4bit_compute_dtype)

bnb_config = BitsAndBytesConfig(
    load_in_4bit=use_4bit,
    bnb_4bit_quant_type=bnb_4bit_quant_type,
    bnb_4bit_compute_dtype=compute_dtype,
    bnb_4bit_use_double_quant=use_nested_quant,
)

# Check GPU compatibility with bfloat16
if compute_dtype == torch.float16 and use_4bit:
    major, _ = torch.cuda.get_device_capability()
    if major >= 8:
        print("=" * 80)
        print("Your GPU supports bfloat16: accelerate training with bf16=True")
        print("=" * 80)

# Load base model
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    quantization_config=bnb_config,
    device_map=device_map
)
model.config.use_cache = False
model.config.pretraining_tp = 1

# Load LLaMA tokenizer
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
tokenizer.pad_token = tokenizer.eos_token
tokenizer.padding_side = "right" # Fix weird overflow issue with fp16 training

# Load LoRA configuration
peft_config = LoraConfig(
    lora_alpha=lora_alpha,
    lora_dropout=lora_dropout,
    r=lora_r,
    bias="none",
    task_type="CAUSAL_LM",
)

# Set training parameters
training_arguments = TrainingArguments(
    output_dir=output_dir,
    num_train_epochs=num_train_epochs,
    per_device_train_batch_size=per_device_train_batch_size,
    gradient_accumulation_steps=gradient_accumulation_steps,
    optim=optim,
    save_steps=save_steps,
    logging_steps=logging_steps,
    learning_rate=learning_rate,
    weight_decay=weight_decay,
    fp16=fp16,
    bf16=bf16,
    max_grad_norm=max_grad_norm,
    max_steps=max_steps,
    warmup_ratio=warmup_ratio,
    group_by_length=group_by_length,
    lr_scheduler_type=lr_scheduler_type,
    report_to="wandb"
)

# Set supervised fine-tuning parameters
trainer = SFTTrainer(
    model=model,
    train_dataset=dataset,
    peft_config=peft_config,
    dataset_text_field="text",
    max_seq_length=max_seq_length,
    tokenizer=tokenizer,
    args=training_arguments,
    packing=packing,
)

# Train model
trainer.train()

# Save trained model
trainer.model.save_pretrained(new_model)

pip install wandb

wandb login

docs = []
loader = PyPDFLoader("/content/main_notes.pdf")
docs.extend(loader.load())
text_splitter = RecursiveCharacterTextSplitter(
    chunk_size = 1500,
    chunk_overlap = 150,
    separators=["\n\n", "\n", " ", ""]
)

splits = text_splitter.split_documents(docs)
client = weaviate.Client(
  embedded_options = EmbeddedOptions()
)

vectorstore = Weaviate.from_documents(
    client = client,
    documents = splits,
    embedding = HuggingFaceEmbeddings(model_name = "NousResearch/Llama-2-7b-chat-hf"),
    by_text = False
)
retriever = vectorstore.as_retriever()

template = """You serve as a assistant specialized in answering questions. Utilize the provided context fragments to respond to the question. If unsure, simply state your lack of knowledge. Provide a concise answer within three sentences at most.
Question: {question}
Context: {context}
Answer:
"""
prompt = ChatPromptTemplate.from_template(template)

llm = HuggingFaceHub(
    repo_id=repo_id, model_kwargs={"temperature": 0.5, "max_length": 64}
)
rag_chain = (
    {"context": retriever,  "question": RunnablePassthrough()}
    | prompt
    | llm
    | StrOutputParser()
)

query = "Tell me about Support Vector Machine"
rag_chain.invoke(query)