Update app.py

app.py

@@ -0,0 +1,98 @@
+import os
+import argparse
+from datasets import load_dataset
+from langchain.schema import Document
+from langchain.vectorstores import Chroma
+from langchain.embeddings import HuggingFaceEmbeddings
+from langchain.llms import LlamaCpp
+from langchain.chains import RetrievalQA
+from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler
+
+# Initialize the database
+def initialize_database():
+    print("🔹 Loading medical dataset...")
+    ds = load_dataset("lavita/ChatDoctor-HealthCareMagic-100k", split="train")
+    qa_pairs = [{"question": x["instruction"], "answer": x["output"]} for x in ds.select(range(1000))]
+
+    # Convert to LangChain Documents
+    print("🔹 Converting to LangChain documents...")
+    docs = [
+        Document(
+            page_content=f"Question: {item['question']}\nAnswer: {item['answer']}",
+            metadata={"source": "ChatDoctor"}
+        )
+        for item in qa_pairs
+    ]
+    
+    # Embedding documents
+    print("🔹 Embedding documents...")
+    embedding_model = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
+
+    # ChromaDB setup
+    persist_dir = "./chroma_medical_db"
+    if not os.path.exists(persist_dir):
+        print("🔹 Creating new ChromaDB...")
+        vectorstore = Chroma.from_documents(docs, embedding_model, persist_directory=persist_dir)
+        vectorstore.persist()
+    else:
+        print("🔹 Loading existing ChromaDB...")
+        vectorstore = Chroma(persist_directory=persist_dir, embedding_function=embedding_model)
+
+    # Setup the retriever
+    retriever = vectorstore.as_retriever(search_kwargs={"k": 3})
+
+    # Local LLM setup
+    print("🔹 Loading local LLM model...")
+    llm = LlamaCpp(
+        model_path="models/tinyllama-1.1b-chat-v1.0.Q4_K_M.gguf",
+        n_ctx=1024,
+        temperature=0.7,
+        max_tokens=512,
+        streaming=True,
+        callbacks=[StreamingStdOutCallbackHandler()],
+        verbose=True,
+        f16_kv=True,
+        use_mlock=True,
+        use_mmap=True,
+        n_threads=4,
+        n_batch=64
+    )
+
+    # Build RAG QA chain
+    print("🔹 Building RAG chain...")
+    qa_chain = RetrievalQA.from_chain_type(
+        llm=llm,
+        retriever=retriever,
+        return_source_documents=True
+    )
+
+    return qa_chain
+
+# Function to handle the query
+def handle_query(query):
+    qa_chain = initialize_database()
+    print(f"🔹 Query: {query}")
+    result = qa_chain(query)
+    response = {
+        "answer": result['result'],
+        "sources": result['source_documents']
+    }
+    return response
+
+# Main CLI functionality
+def main():
+    parser = argparse.ArgumentParser(description="Medical Question-Answering CLI Application")
+    parser.add_argument("query", type=str, help="Query to ask the medical AI agent")
+
+    args = parser.parse_args()
+    query = args.query
+
+    result = handle_query(query)
+    print("\n🧠 Answer:")
+    print(result["answer"])
+    print("\nSource Documents:")
+    for doc in result["sources"]:
+        print(doc["text"])
+
+if __name__ == "__main__":
+    main()