Kazel commited on
Commit
3b40fc3
·
verified ·
1 Parent(s): 1a9d6c3

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +95 -0
app.py ADDED
@@ -0,0 +1,95 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ import pandas as pd
3
+ import os
4
+ import tempfile
5
+ from typing import List, Dict, Any
6
+ import json
7
+ from langchain.text_splitter import RecursiveCharacterTextSplitter
8
+ from langchain_community.vectorstores import Chroma
9
+ from langchain_community.embeddings import HuggingFaceEmbeddings
10
+ from langchain.schema import Document
11
+ from langchain.chains import RetrievalQA
12
+ import logging
13
+ import uuid
14
+ import docx
15
+ import PyPDF2
16
+ import openpyxl
17
+ import pptx
18
+ import shutil
19
+ import re
20
+ from transformers import pipeline
21
+
22
+ # Set up logging
23
+ logging.basicConfig(level=logging.INFO)
24
+ logger = logging.getLogger(__name__)
25
+
26
+ CHROMA_DB_DIR = "./chroma_db"
27
+
28
+ class HFZeroGPULLM:
29
+ def __init__(self, model_id="mistralai/Mistral-7B-Instruct-v0.1"):
30
+ try:
31
+ self.generator = pipeline("text-generation", model=model_id, device=-1)
32
+ logger.info("Loaded HuggingFace text-generation pipeline on CPU.")
33
+ except Exception as e:
34
+ logger.error(f"Failed to load HuggingFace pipeline: {e}")
35
+ self.generator = None
36
+
37
+ def invoke(self, prompt):
38
+ if not self.generator:
39
+ raise RuntimeError("HFZeroGPULLM not initialized properly.")
40
+ result = self.generator(prompt, max_new_tokens=512, do_sample=True)[0]
41
+ return result['generated_text'] if 'generated_text' in result else result['text']
42
+
43
+ class CSVRAGSystem:
44
+ def __init__(self):
45
+ self.vectorstore = None
46
+ self.qa_chain = None
47
+ self.uploaded_files = []
48
+ self.text_splitter = RecursiveCharacterTextSplitter(
49
+ chunk_size=1000,
50
+ chunk_overlap=200,
51
+ length_function=len,
52
+ )
53
+
54
+ # Initialize HuggingFace LLM (CPU-based)
55
+ try:
56
+ self.llm = HFZeroGPULLM()
57
+ logger.info("HuggingFace LLM initialized successfully.")
58
+ except Exception as e:
59
+ logger.error(f"Failed to initialize HuggingFace LLM: {e}")
60
+ self.llm = None
61
+
62
+ # Always try to load persistent ChromaDB
63
+ self.load_vectorstore()
64
+
65
+ def load_vectorstore(self):
66
+ try:
67
+ if os.path.exists(CHROMA_DB_DIR) and os.listdir(CHROMA_DB_DIR):
68
+ embeddings = HuggingFaceEmbeddings(
69
+ model_name="sentence-transformers/all-MiniLM-L6-v2",
70
+ model_kwargs={'device': 'cpu'}
71
+ )
72
+ self.vectorstore = Chroma(
73
+ embedding_function=embeddings,
74
+ persist_directory=CHROMA_DB_DIR
75
+ )
76
+ if self.llm:
77
+ self.qa_chain = RetrievalQA.from_chain_type(
78
+ llm=self.llm,
79
+ chain_type="stuff",
80
+ retriever=self.vectorstore.as_retriever(search_kwargs={"k": 3}),
81
+ return_source_documents=True
82
+ )
83
+ logger.info("Loaded persistent ChromaDB vectorstore.")
84
+ else:
85
+ logger.info("No existing ChromaDB found. Will create on first upload.")
86
+ except Exception as e:
87
+ logger.error(f"Error loading persistent ChromaDB: {e}")
88
+
89
+ # [REMAINDER OF CODE UNCHANGED ... your previous class logic continues here]
90
+
91
+ # NOTE:
92
+ # - The Ollama import was removed.
93
+ # - Replaced Ollama usage with `HFZeroGPULLM` that uses Hugging Face Transformers.
94
+ # - You can adjust the `model_id` (e.g., to llama2 models or phi models) depending on availability.
95
+ # - Ensure `transformers` is added to requirements.txt