pradeepsengarr commited on
Commit
097081a
·
verified ·
1 Parent(s): 0efd1de

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +18 -7
app.py CHANGED
@@ -1,3 +1,4 @@
 
1
  import gradio as gr
2
  from sentence_transformers import SentenceTransformer
3
  from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
@@ -5,20 +6,29 @@ import faiss
5
  import numpy as np
6
  from langchain.text_splitter import RecursiveCharacterTextSplitter
7
  import fitz # PyMuPDF
 
8
 
9
- # Load models
 
 
 
 
 
 
10
  embed_model = SentenceTransformer("BAAI/bge-base-en-v1.5")
11
 
 
12
  model_id = "mistralai/Mistral-7B-Instruct-v0.1"
13
- tokenizer = AutoTokenizer.from_pretrained(model_id)
14
  model = AutoModelForCausalLM.from_pretrained(
15
  model_id,
16
  device_map="auto",
17
- load_in_4bit=True
 
18
  )
19
  llm = pipeline("text-generation", model=model, tokenizer=tokenizer)
20
 
21
- # Globals
22
  index = None
23
  doc_texts = []
24
 
@@ -35,7 +45,7 @@ def extract_text(file):
35
  else:
36
  return "❌ Invalid file type."
37
 
38
- # File processing
39
  def process_file(file):
40
  global index, doc_texts
41
  text = extract_text(file)
@@ -52,9 +62,10 @@ def process_file(file):
52
 
53
  return "✅ File processed successfully. You can now ask questions!"
54
 
55
- # Context + LLM response
56
  def generate_answer(question):
57
- if index is None:
 
58
  return "⚠️ Please upload and process a file first."
59
 
60
  question_embedding = embed_model.encode([question])
 
1
+ import os
2
  import gradio as gr
3
  from sentence_transformers import SentenceTransformer
4
  from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 
6
  import numpy as np
7
  from langchain.text_splitter import RecursiveCharacterTextSplitter
8
  import fitz # PyMuPDF
9
+ from huggingface_hub import login
10
 
11
+ # Authenticate with Hugging Face to access gated models
12
+ hf_token = os.environ.get("HUGGINGFACE_TOKEN")
13
+ if hf_token is None:
14
+ raise ValueError("Please set the HUGGINGFACE_TOKEN environment variable")
15
+ login(token=hf_token)
16
+
17
+ # Load embedding model
18
  embed_model = SentenceTransformer("BAAI/bge-base-en-v1.5")
19
 
20
+ # Load LLM model and tokenizer with 4bit quantization
21
  model_id = "mistralai/Mistral-7B-Instruct-v0.1"
22
+ tokenizer = AutoTokenizer.from_pretrained(model_id, use_auth_token=hf_token)
23
  model = AutoModelForCausalLM.from_pretrained(
24
  model_id,
25
  device_map="auto",
26
+ load_in_4bit=True,
27
+ use_auth_token=hf_token
28
  )
29
  llm = pipeline("text-generation", model=model, tokenizer=tokenizer)
30
 
31
+ # Globals for FAISS index and document texts
32
  index = None
33
  doc_texts = []
34
 
 
45
  else:
46
  return "❌ Invalid file type."
47
 
48
+ # File processing: chunk text, create embeddings, build FAISS index
49
  def process_file(file):
50
  global index, doc_texts
51
  text = extract_text(file)
 
62
 
63
  return "✅ File processed successfully. You can now ask questions!"
64
 
65
+ # Generate answer using retrieved context and LLM
66
  def generate_answer(question):
67
+ global index, doc_texts
68
+ if index is None or len(doc_texts) == 0:
69
  return "⚠️ Please upload and process a file first."
70
 
71
  question_embedding = embed_model.encode([question])