Update app.py

app.py

@@ -1,6 +1,5 @@
 import gradio as gr
 from huggingface_hub import InferenceClient
-
 import nltk
 import PyPDF2
 
@@ -13,35 +12,40 @@ nltk.download("punkt", quiet=True)
 # Initialize the Hugging Face model client
 client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")
 
-def respond(message, history, system_message, max_tokens, temperature, top_p, file_content):
+def respond_chunked(message, history, system_message, max_tokens, temperature, top_p, file_content):
     """
-    Calls the Hugging Face model for a response.
-    Includes file content in the system message for context.
+    Calls the Hugging Face model for a response with support for chunked file content.
     """
-    # Append file content to system message if available
-    if file_content and file_content.strip():
-        system_message += f"\n\nFile content:\n{file_content}"
-    
-    # Prepare the message payload
-    messages = [{"role": "system", "content": system_message}]
-    for user, assistant in history:
-        if user:
-            messages.append({"role": "user", "content": user})
-        if assistant:
-            messages.append({"role": "assistant", "content": assistant})
-    messages.append({"role": "user", "content": message})
-    
-    # Get the model response
-    try:
-        completion = client.chat_completion(
-            messages,
-            max_tokens=max_tokens,
-            temperature=temperature,
-            top_p=top_p,
-        )
-        return completion.choices[0].message["content"]
-    except Exception as e:
-        return f"Error: {e}"
+    # Split file content into manageable chunks
+    chunks = chunk_text(file_content, max_chunk_size=1500)
+    combined_response = ""
+
+    # Process each chunk and append to the response
+    for chunk in chunks:
+        # Append chunk to system message for context
+        chunked_system_message = f"{system_message}\n\nFile Content Chunk:\n{chunk}"
+        
+        # Prepare the message payload
+        messages = [{"role": "system", "content": chunked_system_message}]
+        for user, assistant in history:
+            if user:
+                messages.append({"role": "user", "content": user})
+            if assistant:
+                messages.append({"role": "assistant", "content": assistant})
+        messages.append({"role": "user", "content": message})
+        
+        try:
+            completion = client.chat_completion(
+                messages,
+                max_tokens=max_tokens,
+                temperature=temperature,
+                top_p=top_p,
+            )
+            combined_response += completion.choices[0].message["content"] + "\n"
+        except Exception as e:
+            combined_response += f"Error processing chunk: {e}\n"
+
+    return combined_response.strip()
 
 ###############################################################################
 #                        File Upload & Parsing Functions                      #
@@ -79,16 +83,50 @@ def load_files(files):
             combined_text += f"Error processing file {file}: {e}\n"
     return combined_text
 
+###############################################################################
+#                            Chunking Function                                #
+###############################################################################
+
+def chunk_text(text, max_chunk_size=1500):
+    """
+    Splits text into chunks of up to `max_chunk_size` tokens (approximate).
+    """
+    from nltk.tokenize import sent_tokenize
+
+    sentences = sent_tokenize(text)
+    chunks = []
+    current_chunk = ""
+    current_tokens = 0
+
+    def approximate_token_count(text):
+        # Naive tokenization approximation
+        return len(text.split())
+
+    for sentence in sentences:
+        sentence_tokens = approximate_token_count(sentence)
+        if current_tokens + sentence_tokens <= max_chunk_size:
+            current_chunk += " " + sentence
+            current_tokens += sentence_tokens
+        else:
+            if current_chunk:
+                chunks.append(current_chunk.strip())
+            current_chunk = sentence
+            current_tokens = sentence_tokens
+    if current_chunk:
+        chunks.append(current_chunk.strip())
+
+    return chunks
+
 ###############################################################################
 #                             Gradio UI Layout                                #
 ###############################################################################
 
 with gr.Blocks() as demo:
-    gr.Markdown("# **Chat with File Context**")
+    gr.Markdown("# **Chat with File Context (Chunking for Large Files)**")
     gr.Markdown(
         """
-This app lets you upload file(s) and chat with an AI assistant.
-Uploaded file content will provide context for the conversation.
+This app lets you upload large file(s) and chat with an AI assistant.
+Uploaded file content will be processed in chunks to ensure smooth handling.
         """
     )
     
@@ -111,14 +149,14 @@ Uploaded file content will provide context for the conversation.
     temperature = gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature")
     top_p = gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-p")
 
-    # Chat Function
+    # Chat Function with Chunking
     def chat_function(user_message, history, file_content, system_prompt, max_tokens, temperature, top_p):
         if not user_message.strip():
             return "", history
         # Append user's message to the chat history
         history.append((user_message, ""))
-        # Get response from the model
-        assistant_response = respond(user_message, history, system_prompt, max_tokens, temperature, top_p, file_content)
+        # Get response from the model with chunking
+        assistant_response = respond_chunked(user_message, history, system_prompt, max_tokens, temperature, top_p, file_content)
         history[-1] = (user_message, assistant_response)
         return "", history