# Creating a RetrievalQA chain with PySpark-based document loading
number_of_articles = 20
 
wikipedia_dataframe = spark.read.parquet("databricks-datasets/wikipedia-datasets/data-001/en_wikipedia/articles-only-parquet/*").limit(number_of_articles)
display(wikipedia_dataframe)

# The following lines are all that is needed for loading data from a PySpark Dataframe into Langchain
from langchain.document_loaders import PySparkDataFrameLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
 
loader = PySparkDataFrameLoader(spark, wikipedia_dataframe, page_content_column="text")
documents = loader.load()
 
text_splitter = RecursiveCharacterTextSplitter(chunk_size=3000, chunk_overlap=0)
texts = text_splitter.split_documents(documents)
print(f"Number of documents: {len(texts)}")

# Create a FAISS vector store using HuggingfaceEmbeddings
from langchain.embeddings import OpenAIEmbeddings
from langchain.vectorstores import FAISS
 
embeddings = OpenAIEmbeddings()
db = FAISS.from_documents(texts, embeddings)

# Create a RetrievalQA chain
from langchain.chains import RetrievalQA
from langchain import OpenAI
 
retrieval_qa = RetrievalQA.from_chain_type(llm=OpenAI(), chain_type="stuff", retriever=db.as_retriever())

# Query the RetrievalQA chain
query = "Who is Harrison Schmitt"
result = retrieval_qa({"query": query})
print("Result:", result["result"])

# Logging the chain with MLflow
import mlflow
 
persist_directory = "langchain/faiss_index"
db.save_local(persist_directory)
 
def load_retriever(persist_directory):
  embeddings = OpenAIEmbeddings()
  db = FAISS.load_local(persist_directory, embeddings)
  return db.as_retriever()
 
# Log the RetrievalQA chain
with mlflow.start_run() as mlflow_run:
  logged_model = mlflow.langchain.log_model(
    retrieval_qa,
    "retrieval_qa_chain",
    loader_fn=load_retriever,
    persist_dir=persist_directory,
  )
  
 # Loading the chain using MLflow
 model_uri = f"runs:/{ mlflow_run.info.run_id }/retrieval_qa_chain"
 
loaded_pyfunc_model = mlflow.pyfunc.load_model(model_uri)
langchain_input = {"query": "Who is Harrison Schmitt"}
loaded_pyfunc_model.predict([langchain_input])

More Simple

from langchain.chains import LLMChain
from langchain.prompts import PromptTemplate
from langchain.llms import OpenAI

# Initialize the OpenAI LLM (Language Model)
llm = OpenAI(openai_api_key="your_openai_api_key")
# Define the prompt template
prompt_template = PromptTemplate(
    input_variables=["input_text"],
    template="Summarize the following text: {input_text}"
)
# Create the LLMChain
llm_chain = LLMChain(
    llm=llm,
    prompt=prompt_template
)
# Example input
input_text = "Is it true? That LangChain is a framework for building applications powered by language models. It allows developers to create advanced chains of prompts, use different language models, and deploy them efficiently."
# Run the chain with input
output = llm_chain.run(input_text=input_text)
# Display the output
print("Summary:", output)