gmustafa413 commited on
Commit
ddc98da
Β·
verified Β·
1 Parent(s): 5f0bc37

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +63 -31
app.py CHANGED
@@ -24,8 +24,11 @@ class GeminiRAGSystem:
24
  self.dataset_loaded = False
25
  self.gemini_api_key = os.getenv("AIzaSyASrFvE3gFPigihza0JTuALzZmBx0Kc3d0")
26
 
27
- # Initialize embedding model
28
  try:
 
 
 
29
  self.embedding_model = SentenceTransformer(MODEL_NAME)
30
  except Exception as e:
31
  raise RuntimeError(f"Failed to initialize embedding model: {str(e)}")
@@ -35,11 +38,17 @@ class GeminiRAGSystem:
35
  genai.configure(api_key=self.gemini_api_key)
36
 
37
  def load_dataset(self):
38
- """Load dataset from Hugging Face"""
39
  try:
40
  with gr.Progress() as progress:
41
  progress(0.1, desc="πŸ“¦ Downloading dataset...")
42
- dataset = load_dataset(DATASET_NAME, split='train')
 
 
 
 
 
 
43
 
44
  progress(0.5, desc="πŸ”¨ Processing dataset...")
45
  if 'text' in dataset.features:
@@ -50,7 +59,11 @@ class GeminiRAGSystem:
50
  raise ValueError("Dataset must have 'text' or 'context' field")
51
 
52
  progress(0.7, desc="🧠 Creating embeddings...")
53
- embeddings = self.embedding_model.encode(self.chunks, show_progress_bar=False)
 
 
 
 
54
  self.index = faiss.IndexFlatL2(embeddings.shape[1])
55
  self.index.add(embeddings.astype('float32'))
56
 
@@ -58,25 +71,28 @@ class GeminiRAGSystem:
58
  progress(1.0, desc="βœ… Dataset loaded successfully!")
59
  return True
60
  except Exception as e:
61
- gr.Warning(f"Failed to load dataset: {str(e)}")
62
  return False
63
 
64
  def get_relevant_context(self, query: str) -> str:
65
- """Retrieve most relevant chunks"""
66
  if not self.index:
67
  return ""
68
 
69
- query_embed = self.embedding_model.encode([query])
70
- _, indices = self.index.search(query_embed.astype('float32'), k=TOP_K)
71
-
72
- context = []
73
- for idx in indices[0]:
74
- if idx < len(self.chunks):
75
- context.append(self.chunks[idx])
76
- return "\n\n".join(context)
77
-
 
 
 
78
  def generate_response(self, query: str) -> str:
79
- """Generate response using Gemini"""
80
  if not self.dataset_loaded:
81
  return "⚠️ Please load the dataset first"
82
  if not self.gemini_api_key:
@@ -97,39 +113,55 @@ class GeminiRAGSystem:
97
  response = model.generate_content(prompt)
98
  return response.text
99
  except Exception as e:
100
- return f"⚠️ Error: {str(e)}"
101
 
102
- # Initialize system
103
- rag_system = GeminiRAGSystem()
 
 
 
104
 
105
  # Create interface
106
- with gr.Blocks(title="RAG Chatbot") as app:
107
- gr.Markdown("# UE_ChatBot")
108
 
109
  with gr.Row():
110
  with gr.Column():
111
- load_btn = gr.Button("πŸš€ Load Dataset", variant="primary")
112
- status = gr.Markdown("ℹ️ Click to load dataset")
113
 
114
  with gr.Column():
115
- chatbot = gr.Chatbot()
116
  query = gr.Textbox(label="Your question", placeholder="Ask about the dataset...")
117
- submit_btn = gr.Button("πŸ“€ Submit", variant="primary")
 
 
118
 
119
  # Event handlers
120
  def load_dataset():
121
- if rag_system.load_dataset():
122
- return "βœ… Dataset ready! You can now ask questions."
123
- return "❌ Failed to load dataset"
 
 
 
124
 
125
  def respond(message, chat_history):
126
- response = rag_system.generate_response(message)
127
- chat_history.append((message, response))
128
- return "", chat_history
 
 
 
 
 
 
 
129
 
130
  load_btn.click(load_dataset, outputs=status)
131
  submit_btn.click(respond, [query, chatbot], [query, chatbot])
132
  query.submit(respond, [query, chatbot], [query, chatbot])
 
133
 
134
  if __name__ == "__main__":
135
  app.launch(share=True)
 
24
  self.dataset_loaded = False
25
  self.gemini_api_key = os.getenv("AIzaSyASrFvE3gFPigihza0JTuALzZmBx0Kc3d0")
26
 
27
+ # Initialize embedding model with explicit version compatibility
28
  try:
29
+ # Workaround for huggingface_hub compatibility
30
+ import huggingface_hub
31
+ huggingface_hub.__version__ = "0.13.4" # Force compatible version
32
  self.embedding_model = SentenceTransformer(MODEL_NAME)
33
  except Exception as e:
34
  raise RuntimeError(f"Failed to initialize embedding model: {str(e)}")
 
38
  genai.configure(api_key=self.gemini_api_key)
39
 
40
  def load_dataset(self):
41
+ """Load dataset from Hugging Face with compatibility fallbacks"""
42
  try:
43
  with gr.Progress() as progress:
44
  progress(0.1, desc="πŸ“¦ Downloading dataset...")
45
+
46
+ # Workaround for dataset loading
47
+ dataset = load_dataset(
48
+ DATASET_NAME,
49
+ split='train',
50
+ download_config={"use_auth_token": False}
51
+ )
52
 
53
  progress(0.5, desc="πŸ”¨ Processing dataset...")
54
  if 'text' in dataset.features:
 
59
  raise ValueError("Dataset must have 'text' or 'context' field")
60
 
61
  progress(0.7, desc="🧠 Creating embeddings...")
62
+ embeddings = self.embedding_model.encode(
63
+ self.chunks,
64
+ show_progress_bar=False,
65
+ convert_to_numpy=True
66
+ )
67
  self.index = faiss.IndexFlatL2(embeddings.shape[1])
68
  self.index.add(embeddings.astype('float32'))
69
 
 
71
  progress(1.0, desc="βœ… Dataset loaded successfully!")
72
  return True
73
  except Exception as e:
74
+ gr.Warning(f"Dataset loading error: {str(e)}")
75
  return False
76
 
77
  def get_relevant_context(self, query: str) -> str:
78
+ """Retrieve most relevant chunks with version-safe operations"""
79
  if not self.index:
80
  return ""
81
 
82
+ try:
83
+ query_embed = self.embedding_model.encode(
84
+ [query],
85
+ convert_to_numpy=True
86
+ ).astype('float32')
87
+
88
+ _, indices = self.index.search(query_embed, k=TOP_K)
89
+ return "\n\n".join([self.chunks[i] for i in indices[0] if i < len(self.chunks)])
90
+ except Exception as e:
91
+ print(f"Search error: {str(e)}")
92
+ return ""
93
+
94
  def generate_response(self, query: str) -> str:
95
+ """Generate response with robust error handling"""
96
  if not self.dataset_loaded:
97
  return "⚠️ Please load the dataset first"
98
  if not self.gemini_api_key:
 
113
  response = model.generate_content(prompt)
114
  return response.text
115
  except Exception as e:
116
+ return f"⚠️ API Error: {str(e)}"
117
 
118
+ # Initialize system with compatibility checks
119
+ try:
120
+ rag_system = GeminiRAGSystem()
121
+ except Exception as e:
122
+ raise RuntimeError(f"System initialization failed: {str(e)}")
123
 
124
  # Create interface
125
+ with gr.Blocks(title="UE Chatbot") as app:
126
+ gr.Markdown("UE 24 Hour Service")
127
 
128
  with gr.Row():
129
  with gr.Column():
130
+ load_btn = gr.Button("Load Dataset", variant="primary")
131
+ status = gr.Markdown("System ready - Load dataset to begin")
132
 
133
  with gr.Column():
134
+ chatbot = gr.Chatbot(height=500)
135
  query = gr.Textbox(label="Your question", placeholder="Ask about the dataset...")
136
+ with gr.Row():
137
+ submit_btn = gr.Button("Submit", variant="primary")
138
+ clear_btn = gr.Button("Clear", variant="secondary")
139
 
140
  # Event handlers
141
  def load_dataset():
142
+ try:
143
+ if rag_system.load_dataset():
144
+ return "Dataset ready! Ask questions now."
145
+ return "Failed to load dataset"
146
+ except Exception as e:
147
+ return f" Error: {str(e)}"
148
 
149
  def respond(message, chat_history):
150
+ try:
151
+ response = rag_system.generate_response(message)
152
+ chat_history.append((message, response))
153
+ return "", chat_history
154
+ except Exception as e:
155
+ chat_history.append((message, f"Error: {str(e)}"))
156
+ return "", chat_history
157
+
158
+ def clear_chat():
159
+ return []
160
 
161
  load_btn.click(load_dataset, outputs=status)
162
  submit_btn.click(respond, [query, chatbot], [query, chatbot])
163
  query.submit(respond, [query, chatbot], [query, chatbot])
164
+ clear_btn.click(clear_chat, outputs=chatbot)
165
 
166
  if __name__ == "__main__":
167
  app.launch(share=True)