Update app.py

app.py

@@ -3,17 +3,25 @@ from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain_huggingface import HuggingFaceEmbeddings
 #from langchain.vectorstores import FAISS
 from langchain_community.vectorstores import FAISS
-import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
 from langchain_huggingface.llms import HuggingFacePipeline
 from langchain.prompts import PromptTemplate
-from transformers import pipeline
+
 from langchain_core.output_parsers import StrOutputParser
 from langchain_core.runnables import RunnablePassthrough
 import glob
 import gradio as gr
 
-# Prepare the data
+from langchain_community.embeddings import HuggingFaceInferenceAPIEmbeddings
+
+import os
+
+secret_value_hf = os.getenv("hf_token")
+
+hf_embeddings = HuggingFaceInferenceAPIEmbeddings(
+    api_key=secret_value_hf,
+    model_name="sentence-transformers/all-MiniLM-l6-v2"
+)
 
 md_path = glob.glob( "md_files/*.md")
 
@@ -30,62 +38,18 @@ doc_splits = text_splitter.split_documents(docs_list)
 # Create the embeddings + retriever
 
 db = FAISS.from_documents(doc_splits,
-                          HuggingFaceEmbeddings(model_name='BAAI/bge-base-en-v1.5'))
-
-
-
-
-# Load quantized model
-
-model_name = 'HuggingFaceH4/zephyr-7b-beta'
-
-bnb_config = BitsAndBytesConfig(
-    load_in_4bit=True,
-    bnb_4bit_use_double_quant=True,
-    bnb_4bit_quant_type="nf4",
-    bnb_4bit_compute_dtype=torch.bfloat16
-)
-
-model = AutoModelForCausalLM.from_pretrained(model_name, quantization_config=bnb_config)
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-
-# Setup the LLM chain
-
-
-
-text_generation_pipeline = pipeline(
-    model=model,
-    tokenizer=tokenizer,
-    task="text-generation",
-    temperature=0.2,
-    do_sample=True,
-    repetition_penalty=1.1,
-    return_full_text=True,
-    max_new_tokens=512,
-)
-
-llm = HuggingFacePipeline(pipeline=text_generation_pipeline)
-
-# search in vector database
-
+                          hf_embeddings)
 
 
+# prompt
 prompt_template = '''You are an assistant for question-answering tasks. 
-
         Here is the context to use to answer the question:
-
         {context} 
-
         Think carefully about the above context. 
-
         Now, review the user question:
-
         {question}
-
         Provide an answer to this questions using only the above context. 
-
         Use three sentences maximum and keep the answer concise.
-
         Answer:'''
 
 prompt = PromptTemplate(
@@ -94,57 +58,145 @@ prompt = PromptTemplate(
 )
 
 
-llm_chain = prompt | llm | StrOutputParser()
-
-
-
-
-retriever = db.as_retriever()
-
-rag_chain = (
- {"context": retriever, "question": RunnablePassthrough()}
-    | llm_chain
-)
-
-#question = "what is advantage of chalcogenide perovskite?"
-
-
-
-
-
-def get_output(is_RAG:str,questions:str):
-    if is_RAG== "RAG":
-        generation2=rag_chain.invoke(questions)
-        return generation2.content
-    else:
-        generation1=llm_chain.invoke({"context":"", "question": questions})
-        return generation1.content
-
-demo = gr.Interface(
-    fn=get_output,
-    inputs=[
-        gr.Radio(
+# gradio interface
+
+def get_output(model_name:str,is_RAG:str,questions:str):
+    if model_name=="mistralai/Mistral-7B-Instruct-v0.2":
+        #repo_id = "mistralai/Mistral-7B-Instruct-v0.2"
+        llm = HuggingFaceEndpoint(
+            repo_id=model_name,
+            max_length=512,
+            temperature=0.2,
+            huggingfacehub_api_token=secret_value_hf,
+        )
+        llm_chain = prompt | llm | StrOutputParser()
+        retriever = db.as_retriever(
+            search_type="similarity",
+            search_kwargs={'k': 4}
+        )
+        
+        rag_chain = (
+         {"context": retriever, "question": RunnablePassthrough()}
+            | llm_chain
+        )
+        if is_RAG== "RAG":
+            generation2=rag_chain.invoke(questions)
+            return generation2
+        else:
+            generation1=llm_chain.invoke({"context":"", "question": questions})
+            return generation1
+    elif model_name=="meta-llama/Llama-3.2-3B-Instruct":
+        llm = HuggingFaceEndpoint(
+            repo_id=model_name,
+            max_length=512,
+            temperature=0.2,
+            huggingfacehub_api_token=secret_value_hf,
+        )
+        llm_chain = prompt | llm | StrOutputParser()
+        retriever = db.as_retriever()
+        
+        rag_chain = (
+         {"context": retriever, "question": RunnablePassthrough()}
+            | llm_chain
+        )
+        if is_RAG== "RAG":
+            generation2=rag_chain.invoke(questions)
+            return generation2
+        else:
+            generation1=llm_chain.invoke({"context":"", "question": questions})
+            return generation1
+    elif model_name=="Qwen/Qwen2.5-72B-Instruct":
+        llm = HuggingFaceEndpoint(
+            repo_id=model_name,
+            max_length=512,
+            temperature=0.2,
+            huggingfacehub_api_token=secret_value_hf,
+        )
+        llm_chain = prompt | llm | StrOutputParser()
+        retriever = db.as_retriever()
+        
+        rag_chain = (
+         {"context": retriever, "question": RunnablePassthrough()}
+            | llm_chain
+        )
+        if is_RAG== "RAG":
+            generation2=rag_chain.invoke(questions)
+            return generation2
+        else:
+            generation1=llm_chain.invoke({"context":"", "question": questions})
+            return generation1
+
+
+
+# Custom CSS to style the output area
+custom_css = """
+#output_area {
+    background-color: #1e1e1e; /* Dark background */
+    color: #ffffff; /* White text */
+    padding: 10px;
+    border-radius: 5px;
+    border: 1px solid #333333; /* Dark border */
+    margin-top: 10px;
+}
+
+#output_area h3 {
+    color: #ffcc00; /* Yellow title color */
+    margin-bottom: 10px;
+}
+"""
+
+with gr.Blocks(title="Ask Questions on Chalcogenide Perovskites",theme=gr.themes.Ocean(),css=custom_css) as demo:
+    gr.Markdown("""
+    ## Retrieval-Augmented Generation for Chalcogenide Perovskites
+    This space implements Retrieval-Augmented Generation (RAG) using large language models, based on Hui Haolei's work on chalcogenide perovskite papers. You can select different models and choose whether to use RAG to enhance the responses.
+    """)
+    
+    with gr.Row():
+        model_name = gr.Radio(
+            choices=["mistralai/Mistral-7B-Instruct-v0.2", "meta-llama/Llama-3.2-3B-Instruct", "Qwen/Qwen2.5-72B-Instruct"],
+            value="mistralai/Mistral-7B-Instruct-v0.2",
+            label="Model Name",
+            info="Select the model you want to use."
+        )
+    
+    with gr.Row():
+        rag = gr.Radio(
             choices=["RAG", "No RAG"],
-            type="value",
-            value="RAG",  # Set default value to "Model 1"
-            label="RAG or not"
-        ),
-        gr.Textbox(label="Input Questions",info="input questions on chalcogenide perovskites")
-    ],
-    outputs="markdown",
-    title="RAG using llm zephyr-7b-beta, embedding model BAAI/bge-base-en-v1.5, based on chalcogenide perovskite papers",
-    description="""
-    ## ask a question to get answer on chalcogenide perovskite; or click on the examples below.
-    """,
-    examples=[["RAG","what is advantage of BaZrS3?"],
-              ["RAG","what is bandgap of SrHfS3?"],
-              ["RAG","why is chalcogenide perovskite important?"]
-              ]
-)
-
-# Launch the Gradio app
-if __name__ == "__main__":
-    demo.launch(share=False)
-
-
-
+            value="RAG",
+            label="RAG or Not",
+            info="Choose whether to use Retrieval-Augmented Generation."
+        )
+    
+    with gr.Row():
+        question = gr.Textbox(
+            label="Input Question",
+            placeholder="Enter your question about chalcogenide perovskites here...",
+            lines=2  # Increase the number of lines for better input experience
+        )
+    
+    with gr.Row():
+        submit_button = gr.Button("Submit")
+    
+    with gr.Row():
+        output = gr.Textbox(label="Response",
+                            lines=10,  # Increase the number of lines for the output area
+                            elem_id="output_area"  # Assign a custom ID for styling
+                            )
+    submit_button.click(
+        fn=get_output,
+        inputs=[model_name, rag, question],
+        outputs=output
+    )
+
+    gr.Examples(
+        examples=[
+            ["mistralai/Mistral-7B-Instruct-v0.2", "RAG", "What is the advantage of BaZrS3?"],
+            ["mistralai/Mistral-7B-Instruct-v0.2", "RAG", "What is the bandgap of SrHfS3?"],
+            ["mistralai/Mistral-7B-Instruct-v0.2", "RAG", "Why is chalcogenide perovskite important?"]
+        ],
+        fn=get_output,
+        inputs=[model_name, rag, question],
+        outputs=output
+    )
+
+demo.launch()