|
from ragatouille import RAGPretrainedModel |
|
from langchain_groq import ChatGroq |
|
from langchain.chains import RetrievalQA |
|
from langchain.memory import ConversationBufferMemory |
|
from langchain.prompts import PromptTemplate |
|
from dotenv import load_dotenv |
|
import os |
|
import streamlit as st |
|
import asyncio |
|
|
|
load_dotenv() |
|
GROQ_API_KEY = os.getenv('GROQ_API_KEY') |
|
|
|
llm = ChatGroq(temperature=0, groq_api_key=GROQ_API_KEY, model_name="llama3-70b-8192") |
|
RAG = RAGPretrainedModel.from_pretrained("colbert-ir/colbertv2.0") |
|
system_prompt = """You are a helpful assistant, you will use the provided context to answer user questions. |
|
Read the given context before answering questions and think step by step. If you can not answer a user question based on |
|
the provided context, inform the user. Do not use any other information for answering user. Provide a detailed answer to the question.""" |
|
prompt_template = ( |
|
system_prompt |
|
+ """ |
|
|
|
Context: {history} \n {context} |
|
User: {question} |
|
Answer:""" |
|
) |
|
prompt = PromptTemplate(input_variables=["history", "context", "question"], template=prompt_template) |
|
memory = ConversationBufferMemory(input_key="question", memory_key="history") |
|
|
|
|
|
def rag(full_string): |
|
|
|
RAG.index( |
|
collection=[full_string], |
|
index_name="vector_db", |
|
max_document_length=512, |
|
split_documents=True, |
|
|
|
) |
|
retriever = RAG.as_langchain_retriever(k=5) |
|
qa = RetrievalQA.from_chain_type( |
|
llm=llm, |
|
chain_type="stuff", |
|
retriever=retriever, |
|
return_source_documents=True, |
|
chain_type_kwargs={"prompt": prompt, "memory": memory}, |
|
) |
|
return qa |
