Updated app.py

app.py

@@ -32,7 +32,6 @@ import torch
 
 db_name = "vector_db"
 folder = "my-knowledge-base/"
-MODEL_NAME = "mistralai/Mistral-7B-Instruct-v0.3"
 
 load_dotenv()
 HF_TOKEN = os.getenv("HF_TOKEN")
@@ -40,13 +39,6 @@ if HF_TOKEN is None:
     raise ValueError("HF_TOKEN is not set. Check your .env file.")
 login(HF_TOKEN, add_to_git_credential=True)
 
-tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, token=HF_TOKEN)
-model = AutoModelForCausalLM.from_pretrained(
-    MODEL_NAME, 
-    device_map="cpu",
-    token=HF_TOKEN
-)
-
 def process_files(files):
     os.makedirs(folder, exist_ok=True)
 
@@ -88,14 +80,7 @@ def process_files(files):
     result = collection.get(include=['embeddings', 'documents', 'metadatas'])
 
     # 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
-    )
+    hf_pipeline = pipeline("document-question-answering", model="impira/layoutlm-document-qa")
 
     llm = HuggingFacePipeline(pipeline=hf_pipeline)
     memory = ConversationBufferMemory(memory_key='chat_history', return_messages=True)