generator_llm_pinecone.py

import torch
from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
from transformers import StoppingCriteriaList, StoppingCriteria
import warnings
from ragas import evaluate
from datasets import Dataset
from ragas.metrics import (
    faithfulness,
    answer_relevancy,
    context_recall,
    context_precision,
)
import os
import pandas as pd
from sentence_transformers import SentenceTransformer
from pinecone import Pinecone
import warnings

warnings.filterwarnings("ignore", category=UserWarning)

# Initialize Variables
# model_name = "AI-Sweden-Models/gpt-sw3-126m-instruct"
model_name = "AI-Sweden-Models/gpt-sw3-1.3b-instruct"
# model_name = "AI-Sweden-Models/gpt-sw3-6.7b-v2-instruct"


with open("language_model\pinecone\key.txt", "r") as f:
    os.environ["OPENAI_API_KEY"] = f.read().strip()

device = "cuda:0" if torch.cuda.is_available() else "cpu"

# Initialize Tokenizer & Model
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)
model.eval()
model.to(device)

document_encoder_model = SentenceTransformer("KBLab/sentence-bert-swedish-cased")


def read_file(file_path: str) -> str:
    """Read the contents of a file."""
    with open(file_path, "r") as file:
        return file.read()


# Note: 'index1' has been pre-created in the pinecone console
# read the pinecone api key from a file
pinecone_api_key = read_file("language_model\pinecone\pincecone_api_key.txt")
pc = Pinecone(api_key=pinecone_api_key)
index = pc.Index("index1")


questions = [
    "Har ett barn rätt att få veta ifall det tillkommit genom insemination eller befruktning utanför kroppen som har utförts med andra könsceller än föräldrarnas egna?",
    "Vem ska anses vara förälder till barnet om modern är gift med en kvinna?",
    "Hur går det till ifall en man vill avsäga sig faderskapet till ett barn som han fått på grund av att han var gift med kvinnan som födde barnet?",
    "Vem får adoptera i Sverige?",
    "Får ett gift par adoptera på egen hand?",
    "Får man adoptera ett barn som fyllt 12 år utan barnets samtycke?",
    "Måste man försörja sina barn livet ut?",
    "Om en TV-reparatör har sönder min Tv-apparat, blir denna då ersättningsskyldig för detta?",
    "Vem äger rätten till ett fotografi?",
    "Hur mycket lagstadgat semester har man i Sverige?",
    "Min dotter är 19 år och hon går i gymnasieskolan. Jag är förälder men inte vårdnadshavare för henne. Hon bor inte hos mig utan hos sin mamma. Är jag fortfarande underhållsskyldig?",
    "Min son är adopterad och har börjat undra varför han inte är lik sina syskon, borde jag förklara för honom att han är adopterad? När och hur borde jag göra det?",
    "Min före detta man som efter vår skilsmässa inte längre har vårdnanden för våra barn kräver att jag ska betala för resorna bär barnen besöker honom. Kan detta verkligen stämma?  Det handlar inte om kostnaden från min sida, utan att det känns som han gör allt för att bråka med mig och pallar inte vika mig hela tiden för honom!",
    "Mina barn är så irriterande, får jag lov att slå dem som straff?",
    "Om jag vill göra en rymdresa med mina barn till månen, måste jag då be om min mans tillåtelse?",
]

ground_truths = [
    "Ja. Ett barn som har tillkommit genom en sådan insemination eller befruktning utanför kroppen som har utförts med andra könsceller än föräldrarnas egna har rätt att av sina föräldrar få veta det. Föräldrarna ska så snart det är lämpligt upplysa barnet om att han eller hon har tillkommit genom en sådan behandling.",
    "Om modern är gift med en kvinna eller är registrerad partner vid barnets födelse, ska hennes make eller registrerade partner anses som barnets förälder.",
    "Vill en man som på grund av äktenskap ska anses som far till ett barn att rätten ska förklara att han inte är far till barnet, ska han väcka talan om detta mot barnet eller, om barnet har avlidit, barnets arvingar.",
    "Den som har fyllt 18 år får adoptera.",
    "Nej, makar och sambor får endast adoptera gemensamt.",
    "Nej, den som har fyllt 12 år får adopteras endast om han eller hon samtycker till adoptionen.",
    "Nej, underhållsskyldigheten upphör när barnet fyller arton år. Går barnet i skolan efter denna tidpunkt, är föräldrarna underhållsskyldiga under den tid som skolgången pågår, dock längst intill dess barnet fyller tjugoett år. Till skolgång räknas studier i grundskolan eller gymnasieskolan och annan jämförlig grundutbildning.",
    "Jag har ingen information om detta.",
    "Den som har skapat ett litterärt eller konstnärligt verk har upphovsrätt till verket, exempel på verk är fotografiskt verk eller något annat alster av bildkonst.",
    "Enligt semesterlagen har alla anställda rätt till 25 dagars semester under ett år.",
    "Ja, du är underhållsskyldig. Föräldrarna är underhållsskyldiga under den tid som skolgången pågår (till skolgång räknas studier i grundskolan eller gymnasieskolan och annan jämförlig grundutbildning), dock längst intill dess barnet fyller tjugoett år. Detta gäller oavsett om föräldern har vårdnaden om barnet eller inte, men hänsyn tas till förälderns ekonomiska förmåga samt barnets egna inkomster.",
    "Ett barn som är adopterat har rätt att av sina föräldrar få veta det. Föräldrarna ska så snart det är lämpligt upplysa barnet om att han eller hon är adopterad.",
    "Ja, det stämmer. Om barnet bor tillsammans med endast en förälder, skall den föräldern ta del i kostnaderna för de resor som föranleds av barnets behov av umgänge med den andra föräldern. Det skall ske efter vad som är skäligt med hänsyn till föräldrarnas ekonomiska förmåga och övriga omständigheter.",
    "Nej, barnaga har varit förbjudet i Sverige sen 1979.",
    "Ja, enligt lagen måste båda föräldrarna vara överrens innan rymdresor med barnen görs.",
]


def query_pincecone_namespace(
    vector_databse_index: Pinecone, q_embedding: str, namespace: str
) -> str:
    result = vector_databse_index.query(
        namespace=namespace,
        vector=q_embedding.tolist(),
        top_k=1,
        include_values=True,
        include_metadata=True,
    )
    results = []
    for match in result.matches:
        results.append(match.metadata["paragraph"])
    return results[0]


def generate_prompt(prompt: str) -> str:
    """Generates a prompt for the GPT-3 model"""
    start_token = "<|endoftext|><s>"
    end_token = "<s>"
    return f"{start_token}\nUser:\n{prompt}\n{end_token}\nBot:\n".strip()


def encode_query(query: str) -> torch.Tensor:
    """Encode the query using the model's tokenizer"""
    return document_encoder_model.encode(query)


class StopOnTokenCriteria(StoppingCriteria):
    def __init__(self, stop_token_id):
        self.stop_token_id = stop_token_id

    def __call__(self, input_ids, scores, **kwargs):
        return input_ids[0, -1] == self.stop_token_id


stop_on_token_criteria = StopOnTokenCriteria(stop_token_id=tokenizer.bos_token_id)
answers = []
contexts = []


# på ett sakligt, kortfattat och formellt vis. Svara
def run_rag_system(with_rag=True):
    for question in questions:
        if with_rag:
            query = query_pincecone_namespace(
                vector_databse_index=index,
                q_embedding=encode_query(query=question),
                namespace="ns-parent-balk",
            )
            prompt = (
                "Besvara följande fråga med max tre meningar och använd och referera till den kontext med lagar du fått nedan: \n"
                + "Fråga: \n"
                + question
                + "\n"
                + "Kontext: \n"
                + query
            )
        else:
            prompt = (
                "Besvara följande fråga med max tre meningar. \n"
                + "Fråga: \n"
                + question
            )

        prompt = generate_prompt(prompt=prompt)
        print(prompt)
        print("-" * 50)

        # # Convert prompt to tokens
        input_ids = tokenizer(prompt, return_tensors="pt")["input_ids"].to(device)

        # Genqerate tokens based om prompt
        generated_token_ids = model.generate(
            inputs=input_ids,
            max_new_tokens=200,
            do_sample=False,
            temperature=0.3,
            top_p=1,
            stopping_criteria=StoppingCriteriaList([stop_on_token_criteria]),
        )[0]

        # Decode the generated tokens
        print("Generated answer: ")
        print("#" * 50)
        generated_text = tokenizer.decode(generated_token_ids[len(input_ids[0]) : -1])

        if with_rag:
            contexts.append([query.strip()])
        answers.append(generated_text.strip())
        print(generated_text)


run_rag_system(with_rag=True)
data = {
    "question": questions,
    "answer": answers,
    "contexts": contexts,
    "ground_truth": ground_truths,
}

# Convert dict to dataset
dataset = Dataset.from_dict(data)

result = evaluate(
    dataset=dataset,
    metrics=[
        context_precision,
        context_recall,
        faithfulness,
        answer_relevancy,
    ],
)

df = result.to_pandas()
print(df)
df.to_csv("ragas_results_two_documents.csv")
print("Results saved to ragas_results_two_documents.csv")

answers = []
run_rag_system(with_rag=False)
data = {
    "question": questions,
    "answer": answers,
    "ground_truth": ground_truths,
}

# create data frama from dict data
df = pd.DataFrame(data)
df.to_csv("gpt_results_no_rag.csv")