Upload generator_llm_pinecone.py

generator_llm_pinecone.py

@@ -0,0 +1,220 @@
+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 över 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. Eftersom en förälder skall fullgöra sin underhållsskyldighet genom att betala underhållsbidrag till barnet även om föräldern inte har vårdnaden om barnet och inte heller varaktigt bor tillsammans med barnet. Därtill, är en förälder underhållsskyldig även efter arton år fyllda om barnet i skolan ( studier i grundskolan eller gymnasieskolan och annan jämförlig grundutbildning) efter denna tidpunkt,  dock längst intill dess barnet fyller tjugoett år.",
+    "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")