Updated app.py

app.py

@@ -15,6 +15,10 @@ from langchain_chroma import Chroma
 from langchain.memory import ConversationBufferMemory
 from langchain.chains import ConversationalRetrievalChain
 from langchain_ollama import ChatOllama
+from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+from langchain.llms import HuggingFacePipeline
+from langchain.memory import ConversationBufferMemory
+from langchain.chains import ConversationalRetrievalChain
 import numpy as np
 from sklearn.manifold import TSNE
 import plotly.graph_objects as go
@@ -25,7 +29,10 @@ import shutil
 
 db_name = "vector_db"
 folder = "my-knowledge-base/"
-MODEL = "llama3.2:latest"
+
+MODEL_NAME = "mistralai/Mistral-7B-Instruct"  # Example: Mistral-7B
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, device_map="auto")
 
 def process_files(files):
     os.makedirs(folder, exist_ok=True)
@@ -67,7 +74,17 @@ def process_files(files):
     collection = vectorstore._collection
     result = collection.get(include=['embeddings', 'documents', 'metadatas'])
 
-    llm = ChatOllama(temperature=0.7, model=MODEL)
+    # HF Pipeline
+    hf_pipeline = pipeline(
+        "text-generation",
+        model=model,
+        tokenizer=tokenizer,
+        max_new_tokens=512,  # Limit output length
+        temperature=0.7,      # Control creativity
+        repetition_penalty=1.2
+    )
+
+    llm = HuggingFacePipeline(pipeline=hf_pipeline)
     memory = ConversationBufferMemory(memory_key='chat_history', return_messages=True)
     retriever = vectorstore.as_retriever(search_kwargs={"k": 10})
     global conversation_chain