AI Legal Assistant
LoginBlog

Juridiskā RAG asistenta izveide ar Gradio - pamācība no "Biedrs"

Ievads

Sveicināti, mēs esam "Biedrs" - Latvijas tehnoloģiju jaunuzņēmums, kas specializējas mākslīgā intelekta un mašīnmācīšanās risinājumos juridisku konsultāciju sniegšanai. Mūsu komanda strādā pie inovatīviem risinājumiem, kas pārsniedz vienkāršas ChatGPT saskarsmes.

Šajā pamācībā mēs parādīsim, kā izveidot juridisko asistentu, izmantojot mākslīgā intelekta tehnoloģiju, kas pazīstama kā RAG (Retrieval Augmented Generation). RAG apvieno informācijas izgūšanu no dokumentiem ar valodas modeļu ģenerēšanas spējām, radot sistēmu, kas var atbildēt uz jautājumiem, balstoties uz konkrētiem dokumentiem.

Mūsu RAG asistents ļaus jums:

  • Ielādēt juridiskus dokumentus
  • Uzdot jautājumus par šiem dokumentiem
  • Saņemt atbildes, kas balstītas uz dokumentu saturu
  • Izmantot vietēji izvietotos valodas modeļus efektīvai darbībai

Nepieciešamie rīki un bibliotēkas

Lai sekotu šai pamācībai, jums būs nepieciešams Google Colab konts vai lokāla Python vide. Mēs izmantosim šādas bibliotēkas:

python
!pip install gradio langchain langchain_community transformers peft bitsandbytes datasets wandb anthropic openai pdf2image pypdf truera langchain_cohere
!pip install chromadb trl

1. Projekta iestatīšana

Vispirms importēsim visas nepieciešamās bibliotēkas un iestatīsim pamata konfigurāciju:

python
import os
import torch
import gradio as gr
import numpy as np
import pandas as pd
import wandb
import requests
import json
from typing import List, Dict, Any, Tuple, Optional
from pathlib import Path
import chromadb
from chromadb.utils import embedding_functions
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
from sentence_transformers import CrossEncoder
from tqdm.auto import tqdm
import random
import time
import logging

# Iestatām žurnalēšanu
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

# Konstantes
DEFAULT_MODEL = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
EMBEDDING_MODEL = "BAAI/bge-small-en-v1.5"
WANDB_PROJECT = "legal-rag-assistant"

2. LegalAssistant klases izveide

Tagad izveidosim galveno klasi, kas apvienos visas mūsu asistenta funkcijas:

python
class LegalAssistant:
    def __init__(self):
        self.documents_loaded = False
        self.model_loaded = False
        self.tokenizer = None
        self.model = None
        self.embedding_function = None
        self.chroma_client = None
        self.chroma_collection = None
        self.reranker = None
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        self.hf_token = None

Šī klase būs mūsu asistenta pamats. Apskatīsim svarīgākās metodes, kas nepieciešamas tā darbībai.

3. HuggingFace autentifikācijas iestatīšana

Lai piekļūtu dažiem modeļiem HuggingFace, jums var būt nepieciešams API tokens. Iestatīsim metodi, kas to apstrādā:

python
def set_hf_token(self, token):
    """Iestatām HuggingFace tokenu ierobežoto modeļu piekļuvei"""
    self.hf_token = token
    wandb.login(key="jūsu_wandb_atslēga", relogin=True)
    return f"HuggingFace tokens iestatīts veiksmīgi: {token[:5]}..." if token else "Tokens nav norādīts"

Šī ir svarīga funkcionalitāte, jo Latvijas kontekstā bieži vien ir nepieciešama piekļuve specializētiem modeļiem, kas var labāk apstrādāt juridisko terminoloģiju.

4. Dokumentu ielāde un vektoru datubāzes izveide

Viens no RAG svarīgākajiem aspektiem ir spēja efektīvi glabāt un meklēt dokumentus. Tam izmantosim ChromaDB:

python
def load_documents(self, documents_dir):
    """Ielādē dokumentus no mapes un izveido vektoru datubāzi"""
    try:
        start_time = time.time()

        if not os.path.exists(documents_dir):
            return f"Kļūda: Mape {documents_dir} neeksistē"

        # Izveidojam iegulšanas funkciju
        self.embedding_function = embedding_functions.SentenceTransformerEmbeddingFunction(
            model_name=EMBEDDING_MODEL
        )

        # Inicializējam ChromaDB
        self.chroma_client = chromadb.PersistentClient(path="./chroma_db")
        
        try:
            self.chroma_collection = self.chroma_client.get_collection(
                name="legal_documents",
                embedding_function=self.embedding_function
            )
            logger.info(f"Atrasta esošā kolekcija ar {self.chroma_collection.count()} dokumentiem")
        except ValueError:
            self.chroma_collection = self.chroma_client.create_collection(
                name="legal_documents",
                embedding_function=self.embedding_function
            )
            logger.info("Izveidota jauna kolekcija")

        # Ielādējam dokumentus no mapes
        files = list(Path(documents_dir).glob("*.txt"))
        if not files:
            return "Norādītajā mapē nav atrasti .txt faili"

        documents = []
        metadatas = []
        ids = []

        for file_path in tqdm(files, desc="Dokumentu ielāde"):
            try:
                with open(file_path, "r", encoding="utf-8") as f:
                    content = f.read()

                # Sadalām garus dokumentus fragmentos ar pārklāšanos
                chunks = self._chunk_text(content, chunk_size=1000, overlap=200)

                for i, chunk in enumerate(chunks):
                    doc_id = f"{file_path.stem}_{i}"
                    documents.append(chunk)
                    metadatas.append({"source": str(file_path), "chunk": i})
                    ids.append(doc_id)
            except Exception as e:
                logger.error(f"Kļūda apstrādājot failu {file_path}: {e}")

        # Pievienojam dokumentus kolekcijai pakās
        batch_size = 100
        for i in range(0, len(documents), batch_size):
            end_idx = min(i + batch_size, len(documents))
            self.chroma_collection.add(
                documents=documents[i:end_idx],
                metadatas=metadatas[i:end_idx],
                ids=ids[i:end_idx]
            )

        # Inicializējam pārkārtotāju
        self.reranker = CrossEncoder('cross-encoder/ms-marco-MiniLM-L-6-v2')

        self.documents_loaded = True
        elapsed_time = time.time() - start_time
        return f"Veiksmīgi ielādēti {len(documents)} dokumentu fragmenti no {len(files)} failiem {elapsed_time:.2f} sekundēs."
    except Exception as e:
        return f"Kļūda ielādējot dokumentus: {str(e)}"

Šī metode:

  1. Izveido iegulšanas (embedding) funkciju, kas tekstus pārvērš skaitliskos vektoros
  2. Izveido vai atver ChromaDB kolekciju dokumentu glabāšanai
  3. Ielādē dokumentus no norādītās mapes
  4. Sadala dokumentus mazākos fragmentos, lai uzlabotu meklēšanas precizitāti
  5. Saglabā dokumentus vektoru datubāzē

5. Dokumentu sadalīšana fragmentos

Lai efektīvi apstrādātu lielus juridiskos dokumentus, mums tie jāsadala mazākos fragmentos:

python
def _chunk_text(self, text, chunk_size=1000, overlap=200):
    """Sadala tekstu pārklājošos fragmentos"""
    if len(text) <= chunk_size:
        return [text]

    chunks = []
    start = 0
    while start < len(text):
        end = min(start + chunk_size, len(text))
        # Mēģinām atrast punktu, jaunu rindu vai atstarpi, kur sadalīt
        if end < len(text):
            for char in ['.', '\n', ' ']:
                pos = text.rfind(char, start, end)
                if pos != -1:
                    end = pos + 1
                    break

        chunks.append(text[start:end])
        start = end - overlap if end - overlap > start else end

    return chunks

Šī funkcija ir īpaši svarīga, jo Latvijas juridiskajiem dokumentiem bieži ir raksturīgi gari paragrāfi un sarežģīta struktūra.

6. Valodas modeļa ielāde

Tagad iestatīsim valodas modeli, kas ģenerēs atbildes uz jautājumiem:

python
def load_model(self, model_name=DEFAULT_MODEL, use_8bit=True, use_4bit=False):
    """Ielādē modeli inferensei"""
    try:
        start_time = time.time()

        # Konfigurējam kvantizāciju
        if use_4bit:
            bnb_config = BitsAndBytesConfig(
                load_in_4bit=True,
                bnb_4bit_use_double_quant=True,
                bnb_4bit_quant_type="nf4",
                bnb_4bit_compute_dtype=torch.bfloat16
            )
            use_8bit = False
        else:
            bnb_config = None

        # Ielādējam tokenizatoru
        self.tokenizer = AutoTokenizer.from_pretrained(
            model_name,
            token=self.hf_token,
            trust_remote_code=True
        )

        # Ielādējam modeli
        if self.tokenizer.pad_token_id is None:
            self.tokenizer.pad_token = self.tokenizer.eos_token

        self.model = AutoModelForCausalLM.from_pretrained(
            model_name,
            token=self.hf_token,
            torch_dtype=torch.float16 if self.device == "cuda" else torch.float32,
            load_in_8bit=use_8bit and self.device == "cuda",
            quantization_config=bnb_config if bnb_config else None,
            device_map="auto" if self.device == "cuda" else None,
            trust_remote_code=True
        )

        self.model_loaded = True
        elapsed_time = time.time() - start_time

        model_params = sum(p.numel() for p in self.model.parameters()) / 1_000_000
        memory_usage = torch.cuda.max_memory_allocated() / 1024**3 if torch.cuda.is_available() else 0

        return f"Modelis '{model_name}' veiksmīgi ielādēts {elapsed_time:.2f} sekundēs. Parametri: {model_params:.2f}M. GPU atmiņa: {memory_usage:.2f} GB"
    except Exception as e:
        return f"Kļūda ielādējot modeli: {str(e)}"

Šī metode:

  1. Konfigurē kvantizāciju, lai samazinātu modeļa atmiņas prasības (svarīgi Colab vidē)
  2. Ielādē valodas modeļa tokenizatoru un pašu modeli
  3. Optimizē GPU atmiņas izmantošanu
  4. Atgriež informāciju par modeļa ielādes procesu

7. RAG izgūšanas implementācija

RAG sistēmas sirds ir relevanto dokumentu izgūšana. Implementēsim šo funkcionalitāti:

python
def rag_retrieval(self, query, k=5, use_reranking=True, analyze_with_wandb=False):
    """Iegūst atbilstošus dokumentus, izmantojot RAG"""
    if not self.documents_loaded:
        return [], "Dokumenti nav ielādēti. Lūdzu, vispirms ielādējiet dokumentus."

    try:
        start_time = time.time()
        # Sākotnējā izgūšana
        initial_k = max(k * 3, 15) if use_reranking else k
        results = self.chroma_collection.query(
            query_texts=[query],
            n_results=initial_k
        )

        documents = results['documents'][0]
        metadatas = results['metadatas'][0]
        distances = results['distances'][0]

        retrieved_docs = [
            {
                "content": doc,
                "metadata": meta,
                "score": 1 - dist/2  # Pārveidojam attālumu līdzības vērtējumā
            }
            for doc, meta, dist in zip(documents, metadatas, distances)
        ]

        # Pielietojam pārkārtošanu, ja iespējota
        if use_reranking and self.reranker:
            pairs = [(query, doc["content"]) for doc in retrieved_docs]
            rerank_scores = self.reranker.predict(pairs)

            for i, score in enumerate(rerank_scores):
                retrieved_docs[i]["rerank_score"] = score

            retrieved_docs.sort(key=lambda x: x["rerank_score"], reverse=True)
            retrieved_docs = retrieved_docs[:k]

        # Reģistrējam izgūšanas metriku W&B, ja iespējots
        if analyze_with_wandb:
            try:
                wandb.init(project=RAG_WANDB_PROJECT, name=f"rag_query_{int(time.time())}")

                # Žurnalējam vaicājumu un izgūšanas statistiku
                wandb.log({
                    "query": query,
                    "num_results": len(retrieved_docs),
                    "retrieval_time": time.time() - start_time,
                    "top_score": retrieved_docs[0]["rerank_score"] if use_reranking else retrieved_docs[0]["score"],
                    "avg_score": sum(doc["rerank_score"] if use_reranking else doc["score"] for doc in retrieved_docs) / len(retrieved_docs),
                    "use_reranking": use_reranking,
                })

                # Izveidojam dataframe izgūto dokumentu vizualizēšanai
                df_data = []
                for i, doc in enumerate(retrieved_docs):
                    df_data.append({
                        "rank": i + 1,
                        "content": doc["content"][:100] + "...",  # Saīsinām attēlošanai
                        "source": doc["metadata"]["source"],
                        "initial_score": doc["score"],
                        "rerank_score": doc.get("rerank_score", None)
                    })

                results_table = wandb.Table(dataframe=pd.DataFrame(df_data))
                wandb.log({"retrieval_results": results_table})
                wandb.finish()
            except Exception as e:
                logger.error(f"Kļūda žurnalējot W&B: {e}")

        elapsed_time = time.time() - start_time
        info = f"Izgūti {len(retrieved_docs)} dokumenti {elapsed_time:.2f} sekundēs"
        return retrieved_docs, info
    except Exception as e:
        return [], f"Kļūda izgūšanas laikā: {str(e)}"

Pārkārtošana (reranking) ir īpaši svarīga juridiskajos jautājumos, jo tā palīdz atlasīt dokumentus, kas vistiešāk attiecas uz konkrēto jautājumu.

8. Atbildes ģenerēšana

Tagad, kad mums ir iespēja izgūt atbilstošus dokumentus, izveidosim metodi atbildes ģenerēšanai:

python
def generate_answer(self, query, use_rag=True, k=5, use_reranking=True, analyze_with_wandb=False):
    """Ģenerē atbildi uz juridisku jautājumu"""
    if not self.model_loaded:
        return "Modelis nav ielādēts. Lūdzu, vispirms ielādējiet modeli."

    try:
        start_time = time.time()
        context_docs = []
        retrieval_info = ""

        # Veicam RAG izgūšanu, ja iespējota
        if use_rag and self.documents_loaded:
            retrieved_docs, retrieval_info = self.rag_retrieval(
                query, k=k, use_reranking=use_reranking, analyze_with_wandb=analyze_with_wandb
            )
            context_docs = retrieved_docs

        # Veidojam uzdevumu
        system_prompt = "Jūs esat juridiskais asistents. Atbildiet uz jautājumu, balstoties uz sniegtajiem juridiskajiem dokumentiem. Norādiet savus avotus."
        prompt = self._build_prompt(system_prompt, query, context_docs)

        # Ģenerējam atbildi
        inputs = self.tokenizer(prompt, return_tensors="pt").to(self.device)

        with torch.no_grad():
            outputs = self.model.generate(
                inputs.input_ids,
                max_new_tokens=512,
                temperature=0.7,
                top_p=0.9,
                do_sample=True
            )

        response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
        # Izvelkam tikai asistenta atbildi
        if "assistant" in response.lower():
            response = response.split("assistant", 1)[1]

        elapsed_time = time.time() - start_time

        # Formatējam avotus citēšanai
        sources = []
        for doc in context_docs:
            source_file = os.path.basename(doc["metadata"]["source"])
            if source_file not in sources:
                sources.append(source_file)

        sources_text = "\n\nAvoti:\n" + "\n".join(sources) if sources else ""

        # Reģistrējam ģenerēšanas metriku W&B, ja iespējots
        if analyze_with_wandb:
            try:
                if not wandb.run:
                    wandb.init(project=RAG_WANDB_PROJECT, name=f"generation_{int(time.time())}")

                wandb.log({
                    "query": query,
                    "response_length": len(response),
                    "generation_time": elapsed_time,
                    "num_sources": len(sources),
                    "used_rag": use_rag,
                    "used_reranking": use_reranking
                })

                wandb.finish()
            except Exception as e:
                logger.error(f"Kļūda žurnalējot W&B: {e}")

        return f"{response}{sources_text}\n\nĢenerēts {elapsed_time:.2f} sekundēs."
    except Exception as e:
        return f"Kļūda ģenerējot atbildi: {str(e)}"

9. Pilna uzdevuma veidošana modelim

Lai efektīvi izmantotu valodas modeli, mums jāizveido pareizi formatēts uzdevums:

python
def _build_prompt(self, system, query, context_docs):
    """Izveido uzdevumu ģenerēšanai"""
    formatted_docs = ""
    if context_docs:
        for i, doc in enumerate(context_docs):
            formatted_docs += f"\nDokuments {i+1} (Avots: {os.path.basename(doc['metadata']['source'])}):\n{doc['content']}\n"

    # Formatējam atkarībā no modeļa tipa
    # Šī ir vienkārša veidne, ko var pielāgot atkarībā no modeļa
    prompt = f"<s>[INST] {system}\n\nJautājums: {query}"

    if formatted_docs:
        prompt += f"\n\nLūk, daži atbilstoši juridiskie dokumenti, kas var palīdzēt:\n{formatted_docs}"

    prompt += " [/INST]"
    return prompt

10. Gradio saskarnes izveide

Visbeidzot, izveidosim grafisku lietotāja saskarni ar Gradio:

python
def create_gradio_interface():
    """Izveido Gradio saskarni juridiskajam asistentam"""
    assistant = LegalAssistant()

    # Definējam Gradio saskarni
    with gr.Blocks(title="Juridiskais RAG Asistents") as app:
        gr.Markdown("# 🏛️ Juridiskais RAG Asistents")
        gr.Markdown("Ielādējiet juridiskus dokumentus, uzdodiet jautājumus un uzlabojiet atbildes.")

        with gr.Tab("Iestatīšana"):
            gr.Markdown("## 1. solis: Konfigurējiet autentifikāciju (pēc izvēles)")
            with gr.Row():
                hf_token = gr.Textbox(type="password", label="HuggingFace tokens (ierobežotiem modeļiem)",
                                    placeholder="Ievadiet savu HuggingFace tokenu", lines=1)
                hf_button = gr.Button("Iestatīt tokenu")
            hf_output = gr.Textbox(label="Tokena statuss", interactive=False)

            gr.Markdown("## 2. solis: Ielādējiet dokumentus")
            with gr.Row():
                docs_dir = gr.Textbox(label="Ceļš uz dokumentu mapi", placeholder="/ceļš/uz/juridiskiem/dokumentiem", value="./pdf_laws")
                load_docs_button = gr.Button("Ielādēt dokumentus")
            docs_output = gr.Textbox(label="Ielādes statuss", interactive=False)

            gr.Markdown("## 3. solis: Ielādējiet modeli")
            with gr.Row():
                model_name = gr.Textbox(label="Modeļa nosaukums vai ceļš", value=DEFAULT_MODEL)
                quant_type = gr.Radio(choices=["8-bit", "4-bit", "None"], value="8-bit", label="Kvantizācija")
                load_model_button = gr.Button("Ielādēt modeli")
            model_output = gr.Textbox(label="Modeļa statuss", interactive=False)

        with gr.Tab("Jautājumu uzdošana"):
            gr.Markdown("## Uzdot juridiskus jautājumus")
            with gr.Row():
                question = gr.Textbox(label="Juridisks jautājums", placeholder="Ievadiet savu juridisko jautājumu šeit", lines=3)

            with gr.Row():
                use_rag = gr.Checkbox(label="Izmantot RAG", value=True)
                use_reranking = gr.Checkbox(label="Izmantot pārkārtošanu", value=True)
                top_k = gr.Slider(minimum=1, maximum=20, value=5, step=1, label="Dokumentu skaits")
                analyze_rag = gr.Checkbox(label="Analizēt RAG ar W&B", value=False)

            answer_button = gr.Button("Saņemt atbildi")
            answer_output = gr.Textbox(label="Atbilde", interactive=False, lines=15)

        # Savieno komponentes
        hf_button.click(assistant.set_hf_token, inputs=hf_token, outputs=hf_output)
        load_docs_button.click(assistant.load_documents, inputs=docs_dir, outputs=docs_output)

        load_model_button.click(
            lambda model, quant: assistant.load_model(
                model_name=model,
                use_8bit=quant == "8-bit",
                use_4bit=quant == "4-bit"
            ),
            inputs=[model_name, quant_type],
            outputs=model_output
        )

        answer_button.click(
            assistant.generate_answer,
            inputs=[question, use_rag, top_k, use_reranking, analyze_rag],
            outputs=answer_output
        )

    return app

if __name__ == "__main__":
    app = create_gradio_interface()
    app.launch(share=True)

11. Lietotnes palaišana Colab vidē

Lai palaistu šo lietotni Google Colab, pievienojiet šādu kodu jūsu Colab piezīmju grāmatiņas beigās:

python
# Palaižam lietotni
app = create_gradio_interface()
app.launch(share=True)

Kad palaižat šo kodu, Gradio izveidos publiski pieejamu URL, kuram varēsiet piekļūt, lai izmēģinātu jūsu juridisko asistentu.

Secinājumi

Šajā pamācībā mēs izveidojām spēcīgu juridisko asistentu, kas izmanto RAG tehnoloģiju, lai atbildētu uz jautājumiem, balstoties uz juridiskiem dokumentiem. Šāda veida risinājumi var būt īpaši vērtīgi Latvijas juridiskajā sektorā, kur automātiska dokumentu analīze var ietaupīt daudz laika un resursu.

"Biedrs" komanda turpinās dalīties ar zināšanām par ģeneratīvo mākslīgo intelektu un tā praktisko pielietojumu. Mēs ticam, ka, izmantojot šādas tehnoloģijas, varam padarīt juridisko konsultāciju pakalpojumus pieejamākus un efektīvākus Latvijā.

Ja jums ir jautājumi vai vēlaties uzzināt vairāk par mūsu risinājumiem, lūdzu, sazinieties ar mums!

Noderīgas saites un resursi

  • HuggingFace modeļi
  • Weights & Biases dokumentācija
  • Gradio dokumentācija
  • ChromaDB dokumentācija

Atslēgas vārdi

mākslīgais intelekts, juridiskais asistents, RAG, Retrieval Augmented Generation, Latvijas tehnoloģijas, dokumentu analīze, juridiskā informācija, tiesību tehnoloģijas, LegalTech, Latvija AI, mašīnmācīšanās, NLP, dabiskās valodas apstrāde, Gradio, HuggingFace, Weights & Biases, ChromaDB, jaunuzņēmums