Reduce document size

app.py

@@ -32,10 +32,15 @@ class CustomRAG:
 
 def answer_question(query):
     docs = load_raw_dataset()
+    llm = get_llm("google/flan-t5-base")
+    vector_database = create_vector_database(
+        docs, "sentence-transformers/all-MiniLM-L6-v2"
+    )
+    prompt_template = get_prompt_template()
     rag = CustomRAG(
-        create_vector_database(docs, "all-MiniLM-L6-v2"),
-        get_llm("google/flan-t5-base"),
-        get_prompt_template(),
+        vector_database,
+        llm,
+        prompt_template,
     )
     response, _ = rag.run(query)
 

indexer.py

@@ -3,7 +3,7 @@ import pandas as pd
 from langchain.schema import Document
 from langchain.embeddings import HuggingFaceEmbeddings
 from langchain.vectorstores import FAISS
-from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
 from langchain.llms import HuggingFacePipeline
 from langchain.prompts import PromptTemplate
 
@@ -15,6 +15,8 @@ def load_raw_dataset():
 
     df["combined"] = df["input"] + " " + df["output"]
 
+    df = df.sample(n=min(5000, len(df)), random_state=42).reset_index(drop=True)
+
     docs = [
         Document(
             page_content=row["combined"],
@@ -34,15 +36,15 @@ def create_vector_database(docs, model_name):
 
 def get_llm(model_name):
     tokenizer = AutoTokenizer.from_pretrained(model_name)
-    model = AutoModelForCausalLM.from_pretrained(
-        model_name, torch_dtype="auto", device_map="auto"
+    model = AutoModelForSeq2SeqLM.from_pretrained(
+        "google/flan-t5-base", torch_dtype="auto", device_map="auto"
     )
 
     pipe = pipeline(
-        "text-generation",
+        "text2text-generation",
         model=model,
         tokenizer=tokenizer,
-        max_new_tokens=300,
+        max_new_tokens=1024,
         temperature=0.7,
         do_sample=True,
     )