Update app.py

app.py

@@ -1,11 +1,87 @@
+import os
+import json
+import pandas as pd
+from transformers import AutoTokenizer, AutoModelForCausalLM
+from langchain_ollama import OllamaLLM, OllamaEmbeddings
+from langchain_community.vectorstores import FAISS
+from langchain_core.prompts import PromptTemplate
+from langchain_core.output_parsers import StrOutputParser
+from operator import itemgetter
 import gradio as gr
-from huggingface_hub import InferenceClient
+from huggingface_hub import login, InferenceClient
 
+USE_HF = True
+MODEL_NAME = "meta-llama/Llama-2-7b-hf"
+
+with open('mini_data.json', 'r', encoding='utf-8') as f:
+    data = json.load(f)
+
+df = pd.DataFrame(data)
+
+documents = [
+    f"Source: {item['Source']}\nApplication: {item['Application']}\nFunction1: {item['Function1']}\nStrategy: {item['Strategy']}"
+    for item in data
+]
+
+if USE_HF:
+    print("Using Hugging Face model...")
+    huggingface_token = os.getenv("AskNature_RAG")
+    login(token=huggingface_token)
+    # Load tokenizer and model from Hugging Face Hub
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, token=huggingface_token)
+    model = AutoModelForCausalLM.from_pretrained(
+        MODEL_NAME,
+        device_map="auto",
+        offload_folder="offload",  # Specify the offload folder
+        token=huggingface_token
+    )
+    embeddings = OllamaEmbeddings(model=MODEL_NAME)
+    lang_model = OllamaLLM(model=MODEL_NAME)
+else:
+    print("Using local Ollama model...")
+    MODEL = "jsk/bio-mistral"
+    embeddings = OllamaEmbeddings(model=MODEL)
+    lang_model = OllamaLLM(model=MODEL)
+
+batch_size = 16
+batched_embeddings = [
+    embeddings.embed_documents(documents[i:i + batch_size])
+    for i in range(0, len(documents), batch_size)
+]
+batched_embeddings = [embed for batch in batched_embeddings for embed in batch]
+
+index_path = "faiss_index"
+if os.path.exists(index_path):
+    vectorstore = FAISS.load_local(index_path, embeddings)
+else:
+    vectorstore = FAISS.from_texts(documents, embeddings)
+    vectorstore.save_local(index_path)
+
+retriever = vectorstore.as_retriever()
+
+template = """
+Answer the question based on the context below. If you can't 
+answer the question, reply "I don't know".
+
+Context: {context}
+
+Question: {question}
 """
-For more information on `huggingface_hub` Inference API support, please check the docs: https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference
-"""
-client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")
+prompt = PromptTemplate.from_template(template)
+
+chain = {
+    "context": itemgetter("question") | retriever,
+    "question": itemgetter("question"),
+} | prompt | lang_model | StrOutputParser()
 
+def rag_qa(question):
+    try:
+        return chain.invoke({'question': question})
+    except Exception as e:
+        return f"Error: {str(e)}"
+
+# Chatbot functionality
+client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")
 
 def respond(
     message,
@@ -39,10 +115,6 @@ def respond(
         response += token
         yield response
 
-
-"""
-For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
-"""
 demo = gr.ChatInterface(
     respond,
     additional_inputs=[
@@ -59,6 +131,5 @@ demo = gr.ChatInterface(
     ],
 )
 
-
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()