Merge pull request #20 from gjyotin305/to_check

Update TigerRag Langchain Folder to Make it Match Newest Movie_Recs Logic #13

TigerRag/demos/langchain/demo.py

@@ -3,70 +3,25 @@
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain.vectorstores import Chroma
 from langchain.embeddings import GPT4AllEmbeddings
-import openai
+from tigerrag.rag.retrieval_augmenters import OpenAIRetrievalAugmenter
+from tigerrag.gar.query_augmenters import OpenAIQueryAugmenter
 import os
 import sys
 
+
 # Please set env var OPENAI_API_KEY for GAR and RAG.
 
 # Sample usage:
 #    python demo.py
 #    python demo.py -number_of_run 4
 
-
-def get_documents_embeddings(documents):
-    # Load documents
-    loader = WebBaseLoader(documents)
-
-    # Split documents
-    text_splitter = RecursiveCharacterTextSplitter(
-        chunk_size=500, chunk_overlap=0)
-    splits = text_splitter.split_documents(loader.load())
-
-    # Embed and store splits
-    vectorstore = Chroma.from_documents(
-        documents=splits, embedding=GPT4AllEmbeddings())
-
-    return vectorstore
-
-
-# EBR
 def ebr(question, vectorstore):
     # Perform similarity search
     docs = vectorstore.similarity_search(question)
 
     return docs[0]
 
 
-# RAG
-def generate_answer_with_rag_gpt3(question, context, openai_text_model):
-    # Retrivel Augmented Generation
-    prompt = f"Context: {context} Question: {question}. Provide a summary or answer:"
-
-    # Generation using GPT-3
-    response = openai.Completion.create(
-        engine=openai_text_model, prompt=prompt, max_tokens=100)
-    answer = response.choices[0].text.strip()
-
-    return answer
-
-
-# GAR
-def generate_answer_with_gar_gpt3(question, context, openai_text_model, vectorstore):
-    # Generation Augmented Retrieval
-    prompt = f"Expand on the query: {question}"
-
-    # Generation using GPT-3
-    response = openai.Completion.create(
-        engine=openai_text_model, prompt=prompt, max_tokens=100)
-    augmented_query = response.choices[0].text.strip()
-
-    # Retrieval
-    answer = ebr(augmented_query, vectorstore)
-
-    return answer
-
-
 def is_intstring(s):
     try:
         int(s)
@@ -98,20 +53,39 @@ def main():
     for index in range(min(num_of_run, len(data["queries"]))):
         question = data["queries"][index]
         # Example usage of EBR
-        vectorstore = get_documents_embeddings(documents)
+        loader = WebBaseLoader(documents)
+        text_splitter = RecursiveCharacterTextSplitter(
+            chunk_size=500, chunk_overlap=0)
+        splits = text_splitter.split_documents(loader.load())
+
+        vectorstore = Chroma.from_documents(
+            documents=splits, embedding=GPT4AllEmbeddings())
         print("The following is EBR output for question: "+question)
         retrieved_context = ebr(question, vectorstore)
         print(retrieved_context)
 
         # Example usage of RAG
-        print("The following is RAG output for question: "+question)
-        print(generate_answer_with_rag_gpt3(
-            question, retrieved_context, 'text-davinci-003'))
 
+        print("The following is RAG output for question: "+question)
+        # Retrivel Augmented Generation
+        prompt_rag = f"""Context: {retrieved_context} Question: {question}.
+        Provide a summary or answer:"""
+        openai_generative_retrieval_augmenter = OpenAIRetrievalAugmenter(
+            "text-davinci-003")
+        answer_rag = openai_generative_retrieval_augmenter.get_augmented_retrieval(
+            prompt_rag)
+        print(answer_rag)
         # Example usage of GAR
+
         print("The following is GAR output for question: "+question)
-        print(generate_answer_with_gar_gpt3(
-            question, retrieved_context, 'text-davinci-003', vectorstore))
+        # print(generate_answer_with_gar_gpt3(
+        #     question, retrieved_context, 'text-davinci-003', vectorstore))
+        prompt_gar = f"Expand on the query: {question}"
+        openai_generative_query_augmenter = OpenAIQueryAugmenter(
+            "text-davinci-003")
+        augmented_query = openai_generative_query_augmenter.get_augmented_query(
+            prompt_gar)
+        print(ebr(augmented_query, vectorstore))
 
 
 if __name__ == "__main__":

TigerRag/tigerrag/base/models.py

@@ -1,9 +1,10 @@
 from enum import Enum
-
+from typing import List
 import numpy as np
 import numpy.typing as npt
 import pandas as pd
 import torch
+
 from transformers import (BertModel, BertTokenizer, RobertaModel,
                           RobertaTokenizer, XLNetModel, XLNetTokenizer)