Copain22 commited on
Commit
8ca6b5f
·
verified ·
1 Parent(s): 9f57aad

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +25 -56
app.py CHANGED
@@ -1,7 +1,7 @@
1
  # 0. Install custom transformers and imports
2
  import os
3
  os.system("pip install git+https://github.com/shumingma/transformers.git")
4
- os.system("pip install sentence-transformers")
5
 
6
  import threading
7
  import torch
@@ -13,13 +13,9 @@ from transformers import (
13
  AutoTokenizer,
14
  TextIteratorStreamer,
15
  )
16
- from sentence_transformers import SentenceTransformer
17
  import gradio as gr
18
  import spaces
19
- import pdfplumber
20
-
21
- from pathlib import Path
22
- from PyPDF2 import PdfReader
23
 
24
  # 1. System prompt
25
  SYSTEM_PROMPT = """
@@ -46,55 +42,28 @@ model = AutoModelForCausalLM.from_pretrained(
46
 
47
  print(f"Model loaded on device: {model.device}")
48
 
49
- # 3. Load PDF files
50
- def load_pdfs(folder_path="."):
51
- docs = []
52
- current_section = None
53
- for pdf_file in Path(folder_path).glob("*.pdf"):
54
- with pdfplumber.open(str(pdf_file)) as pdf:
55
- for page in pdf.pages:
56
- text = page.extract_text()
57
- if text:
58
- lines = text.split("\n")
59
- for line in lines:
60
- line = line.strip()
61
- if not line:
62
- continue
63
-
64
- if line.isupper() and len(line.split()) <= 6:
65
- if current_section:
66
- docs.append(current_section)
67
- current_section = line
68
- else:
69
- if current_section:
70
- current_section += f" | {line}"
71
- else:
72
- current_section = line
73
-
74
- if current_section:
75
- docs.append(current_section)
76
- current_section = None
77
-
78
- return docs
79
 
 
 
80
 
81
- document_chunks = load_pdfs(".")
82
- print(f"Loaded {len(document_chunks)} text chunks from PDFs.")
83
-
84
- # 4. Create embeddings
85
- embedder = SentenceTransformer("all-MiniLM-L6-v2") # Fast small model
86
- doc_embeddings = embedder.encode(document_chunks, normalize_embeddings=True)
87
-
88
- # 5. Retrieval function with float32 fix
89
  def retrieve_context(question, top_k=3):
90
- question_embedding = embedder.encode(question, normalize_embeddings=True)
91
- question_embedding = torch.tensor(question_embedding, dtype=torch.float32)
92
- doc_embeds = torch.tensor(doc_embeddings, dtype=torch.float32)
93
- scores = doc_embeds @ question_embedding
94
- top_indices = torch.topk(scores, k=min(top_k, len(scores))).indices.tolist()
95
- return "\n\n".join([document_chunks[idx] for idx in top_indices])
96
-
97
- # 6. Chat respond function
98
  @spaces.GPU
99
  def respond(
100
  message: str,
@@ -138,11 +107,11 @@ def respond(
138
  response += new_text
139
  yield response
140
 
141
- # 7. Gradio ChatInterface
142
  demo = gr.ChatInterface(
143
  fn=respond,
144
  title="Café Eleven Assistant",
145
- description="Friendly café assistant with real menu knowledge!",
146
  examples=[
147
  [
148
  "What kinds of burgers do you have?",
@@ -152,7 +121,7 @@ demo = gr.ChatInterface(
152
  0.95,
153
  ],
154
  [
155
- "Do you have any gluten-free pastries?",
156
  SYSTEM_PROMPT.strip(),
157
  512,
158
  0.7,
@@ -188,6 +157,6 @@ demo = gr.ChatInterface(
188
  ],
189
  )
190
 
191
- # 8. Launch
192
  if __name__ == "__main__":
193
  demo.launch(share=True)
 
1
  # 0. Install custom transformers and imports
2
  import os
3
  os.system("pip install git+https://github.com/shumingma/transformers.git")
4
+ os.system("pip install python-docx")
5
 
6
  import threading
7
  import torch
 
13
  AutoTokenizer,
14
  TextIteratorStreamer,
15
  )
 
16
  import gradio as gr
17
  import spaces
18
+ from docx import Document
 
 
 
19
 
20
  # 1. System prompt
21
  SYSTEM_PROMPT = """
 
42
 
43
  print(f"Model loaded on device: {model.device}")
44
 
45
+ # 3. Load Menu Text from Word document
46
+ def load_menu_text(docx_path):
47
+ doc = Document(docx_path)
48
+ full_text = []
49
+ for para in doc.paragraphs:
50
+ if para.text.strip():
51
+ full_text.append(para.text.strip())
52
+ return "\n".join(full_text)
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
53
 
54
+ MENU_TEXT = load_menu_text("menu.docx")
55
+ print(f"Loaded menu text from Word document.")
56
 
57
+ # 4. Simple retrieval function (search inside MENU_TEXT)
 
 
 
 
 
 
 
58
  def retrieve_context(question, top_k=3):
59
+ question = question.lower()
60
+ sentences = MENU_TEXT.split("\n")
61
+ matches = [s for s in sentences if any(word in s.lower() for word in question.split())]
62
+ if not matches:
63
+ return "Sorry, I couldn't find relevant menu information."
64
+ return "\n\n".join(matches[:top_k])
65
+
66
+ # 5. Chat respond function
67
  @spaces.GPU
68
  def respond(
69
  message: str,
 
107
  response += new_text
108
  yield response
109
 
110
+ # 6. Gradio ChatInterface
111
  demo = gr.ChatInterface(
112
  fn=respond,
113
  title="Café Eleven Assistant",
114
+ description="Friendly café assistant based on real menu loaded from Word document!",
115
  examples=[
116
  [
117
  "What kinds of burgers do you have?",
 
121
  0.95,
122
  ],
123
  [
124
+ "Do you have gluten-free pastries?",
125
  SYSTEM_PROMPT.strip(),
126
  512,
127
  0.7,
 
157
  ],
158
  )
159
 
160
+ # 7. Launch
161
  if __name__ == "__main__":
162
  demo.launch(share=True)