pradeepsengarr commited on
Commit
378f4da
Β·
verified Β·
1 Parent(s): f1a28c3

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +27 -25
app.py CHANGED
@@ -18,40 +18,42 @@ login(token=hf_token)
18
  # Load embedding model
19
  embed_model = SentenceTransformer("BAAI/bge-base-en-v1.5")
20
 
21
- # Load Mistral without 4bit quantization (CPU-friendly)
22
  model_id = "mistralai/Mistral-7B-Instruct-v0.1"
23
  tokenizer = AutoTokenizer.from_pretrained(model_id, token=hf_token)
24
  model = AutoModelForCausalLM.from_pretrained(
25
  model_id,
26
- device_map={"": "cpu"}, # force CPU
27
- torch_dtype="auto", # safe for CPU
28
  token=hf_token
29
  )
30
  llm = pipeline("text-generation", model=model, tokenizer=tokenizer)
31
 
32
- # Globals
33
  index = None
34
  doc_texts = []
35
 
36
- # Extract text from PDF or TXT
37
- def extract_text(file):
38
  text = ""
39
- file_bytes = file.read()
40
- if file.name.endswith(".pdf"):
41
- pdf_stream = BytesIO(file_bytes)
 
42
  doc = fitz.open(stream=pdf_stream, filetype="pdf")
43
  for page in doc:
44
  text += page.get_text()
45
- elif file.name.endswith(".txt"):
46
- text = file_bytes.decode("utf-8")
 
47
  else:
48
  return "❌ Unsupported file type."
49
  return text
50
 
51
- # Process the file, build FAISS index
52
- def process_file(file):
53
  global index, doc_texts
54
- text = extract_text(file)
55
  if text.startswith("❌"):
56
  return text
57
 
@@ -63,9 +65,9 @@ def process_file(file):
63
  index = faiss.IndexFlatL2(dim)
64
  index.add(embeddings)
65
 
66
- return "βœ… File processed! You can now ask questions."
67
 
68
- # Generate answer using context + LLM
69
  def generate_answer(question):
70
  global index, doc_texts
71
  if index is None or not doc_texts:
@@ -75,30 +77,30 @@ def generate_answer(question):
75
  _, I = index.search(question_emb, k=3)
76
  context = "\n".join([doc_texts[i] for i in I[0]])
77
 
78
- prompt = f"""[System: You are a helpful assistant. Answer based on the context.]
79
 
80
  Context:
81
  {context}
82
 
83
  Question: {question}
84
- Answer:"""
85
 
86
  response = llm(prompt, max_new_tokens=300, do_sample=True, temperature=0.7)
87
  return response[0]["generated_text"].split("Answer:")[-1].strip()
88
 
89
  # Gradio UI
90
- with gr.Blocks(title="RAG Chatbot (CPU Compatible)") as demo:
91
- gr.Markdown("## πŸ“š Upload PDF/TXT and Ask Questions using Mistral-7B")
92
 
93
  with gr.Row():
94
- file_input = gr.File(label="πŸ“ Upload File (.pdf or .txt)", file_types=[".pdf", ".txt"])
95
- upload_status = gr.Textbox(label="Upload Status", interactive=False)
96
 
97
  with gr.Row():
98
- question_box = gr.Textbox(label="❓ Ask a Question", placeholder="What would you like to know?")
99
- answer_box = gr.Textbox(label="πŸ’¬ Answer", interactive=False)
100
 
101
  file_input.change(fn=process_file, inputs=file_input, outputs=upload_status)
102
- question_box.submit(fn=generate_answer, inputs=question_box, outputs=answer_box)
103
 
104
  demo.launch()
 
18
  # Load embedding model
19
  embed_model = SentenceTransformer("BAAI/bge-base-en-v1.5")
20
 
21
+ # Load Mistral LLM (CPU compatible)
22
  model_id = "mistralai/Mistral-7B-Instruct-v0.1"
23
  tokenizer = AutoTokenizer.from_pretrained(model_id, token=hf_token)
24
  model = AutoModelForCausalLM.from_pretrained(
25
  model_id,
26
+ device_map={"": "cpu"}, # Force CPU
27
+ torch_dtype="auto", # Safe for CPU
28
  token=hf_token
29
  )
30
  llm = pipeline("text-generation", model=model, tokenizer=tokenizer)
31
 
32
+ # Global state
33
  index = None
34
  doc_texts = []
35
 
36
+ # Extract text from uploaded file
37
+ def extract_text(file_obj):
38
  text = ""
39
+ file_path = file_obj.name
40
+ if file_path.endswith(".pdf"):
41
+ with open(file_path, "rb") as f:
42
+ pdf_stream = BytesIO(f.read())
43
  doc = fitz.open(stream=pdf_stream, filetype="pdf")
44
  for page in doc:
45
  text += page.get_text()
46
+ elif file_path.endswith(".txt"):
47
+ with open(file_path, "r", encoding="utf-8") as f:
48
+ text = f.read()
49
  else:
50
  return "❌ Unsupported file type."
51
  return text
52
 
53
+ # Process file and build FAISS index
54
+ def process_file(file_obj):
55
  global index, doc_texts
56
+ text = extract_text(file_obj)
57
  if text.startswith("❌"):
58
  return text
59
 
 
65
  index = faiss.IndexFlatL2(dim)
66
  index.add(embeddings)
67
 
68
+ return "βœ… File processed successfully. You can now ask questions!"
69
 
70
+ # Generate answer from FAISS context + LLM
71
  def generate_answer(question):
72
  global index, doc_texts
73
  if index is None or not doc_texts:
 
77
  _, I = index.search(question_emb, k=3)
78
  context = "\n".join([doc_texts[i] for i in I[0]])
79
 
80
+ prompt = f"""<s>[INST] You are a helpful assistant. Use the context below to answer the question.
81
 
82
  Context:
83
  {context}
84
 
85
  Question: {question}
86
+ Answer: [/INST]</s>"""
87
 
88
  response = llm(prompt, max_new_tokens=300, do_sample=True, temperature=0.7)
89
  return response[0]["generated_text"].split("Answer:")[-1].strip()
90
 
91
  # Gradio UI
92
+ with gr.Blocks(title="RAG Chatbot with Mistral-7B (CPU-Friendly)") as demo:
93
+ gr.Markdown("## πŸ€– Upload a PDF/TXT file and ask questions using Mistral-7B")
94
 
95
  with gr.Row():
96
+ file_input = gr.File(label="πŸ“ Upload PDF or TXT", file_types=[".pdf", ".txt"])
97
+ upload_status = gr.Textbox(label="πŸ“₯ Upload Status", interactive=False)
98
 
99
  with gr.Row():
100
+ question_input = gr.Textbox(label="❓ Ask a Question")
101
+ answer_output = gr.Textbox(label="πŸ’¬ Answer", interactive=False)
102
 
103
  file_input.change(fn=process_file, inputs=file_input, outputs=upload_status)
104
+ question_input.submit(fn=generate_answer, inputs=question_input, outputs=answer_output)
105
 
106
  demo.launch()