Update app.py

app.py

@@ -64,37 +64,63 @@ def respond(message, history, system_message, max_tokens, temperature, top_p):
     try:
         model, tokenizer = initialize_model()
         
-        # Get relevant context from the database
+        # Get context from database
         retriever = db.as_retriever(search_kwargs={"k": 2})
         docs = retriever.get_relevant_documents(message)
         context = "\n".join([doc.page_content for doc in docs])
         
-        # Generate the complete prompt
+        # Generate prompt
         prompt = generate_prompt(context=context, question=message, system_prompt=system_message)
         
-        # Set up the streamer
-        streamer = CustomTextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
-        
-        # Set up the pipeline
-        text_pipeline = pipeline(
-            "text-generation",
-            model=model,
-            tokenizer=tokenizer,
+        # Generate response without streamer for direct string output
+        output = text_pipeline(
+            prompt,
             max_new_tokens=max_tokens,
             temperature=temperature,
             top_p=top_p,
             repetition_penalty=1.15,
-            streamer=streamer,
-        )
-        
-        # Generate response
-        _ = text_pipeline(prompt, max_new_tokens=max_tokens)
+            return_full_text=False
+        )[0]['generated_text']
         
-        # Return only the generated response
-        yield streamer.output_text.strip()
+        yield output.strip()
         
     except Exception as e:
         yield f"An error occurred: {str(e)}"
+# def respond(message, history, system_message, max_tokens, temperature, top_p):
+#     try:
+#         model, tokenizer = initialize_model()
+        
+#         # Get relevant context from the database
+#         retriever = db.as_retriever(search_kwargs={"k": 2})
+#         docs = retriever.get_relevant_documents(message)
+#         context = "\n".join([doc.page_content for doc in docs])
+        
+#         # Generate the complete prompt
+#         prompt = generate_prompt(context=context, question=message, system_prompt=system_message)
+        
+#         # Set up the streamer
+#         streamer = CustomTextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
+        
+#         # Set up the pipeline
+#         text_pipeline = pipeline(
+#             "text-generation",
+#             model=model,
+#             tokenizer=tokenizer,
+#             max_new_tokens=max_tokens,
+#             temperature=temperature,
+#             top_p=top_p,
+#             repetition_penalty=1.15,
+#             streamer=streamer,
+#         )
+        
+#         # Generate response
+#         _ = text_pipeline(prompt, max_new_tokens=max_tokens)
+        
+#         # Return only the generated response
+#         yield streamer.output_text.strip()
+        
+#     except Exception as e:
+#         yield f"An error occurred: {str(e)}"
 
 # Create Gradio interface
 demo = gr.ChatInterface(