Update app.py

app.py

@@ -11,142 +11,86 @@ from io import BytesIO
 import docx
 import ntpath
 
-###############################################################################
-# REPLACE WITH YOUR DEFAULT (OPEN-SOURCE) DEEPSEEK API KEY
-###############################################################################
-OPEN_SOURCE_DEEPSEEK_KEY = "YOUR_DEFAULT_API_KEY_HERE"
 
-def chat_with_deepseek(
-    user_message,    # The user’s message to the chatbot
-    history,         # Chat history (handled by Gradio’s ChatInterface)
-    system_message,  # The system message from the left panel
-    user_api_key,    # The user-provided optional HF API key
-    max_tokens,      # The “Max new tokens” slider
-    temperature,     # The “Temperature” slider
-    top_p            # The “Top-p” slider
-):
-    """
-    This function is called every time a user submits a new message.
-    It uses either the user’s HF API key if provided, or a default
-    open-source DeepSeek API key if empty.
-    """
-
-    # Decide which key to use
-    final_api_key = user_api_key.strip() if user_api_key else OPEN_SOURCE_DEEPSEEK_KEY
 
-    # Initialize InferenceClient with the chosen API key
-    client = InferenceClient(token=final_api_key)
+# Initialize clients
+API_KEY = os.environ.get("HF_API_KEY")
+client = InferenceClient(token=API_KEY)
 
-    # Build the prompt or system instruction
-    # You can format your prompt however you like; here’s a simple example:
-    prompt = (
-        f"{system_message.strip()}\n\n"    # System instructions at the top
-        f"User: {user_message}\n"
-        "Assistant:"
-    )
+def process_file(filepath):
+    # Handle different file types and extract text.
+    if not filepath:
+        return ""
+    
+    ext = os.path.splitext(filepath)[1].lower()
+    
+    try:
+        if ext in ['.png', '.jpg', '.jpeg']:
+            image = Image.open(filepath)
+            text = pytesseract.image_to_string(image)
+            return f"IMAGE CONTENT:\n{text}"
+        elif ext == '.pdf':
+            pdf_reader = PyPDF2.PdfReader(filepath)
+            text = "\n".join(page.extract_text() for page in pdf_reader.pages)
+            return f"PDF CONTENT:\n{text}"
+        elif ext == '.docx':
+            doc = docx.Document(filepath)
+            text = "\n".join(para.text for para in doc.paragraphs)
+            return f"DOCUMENT CONTENT:\n{text}"
+        else:
+            return "Unsupported file type"
+    except Exception as e:
+        print(f"File processing error: {e}")
+        return "Error reading file"
+
+def chat(message, history, filepath):
+    file_content = process_file(filepath) if filepath else ""
+    
+    full_prompt = f"""
+    {file_content}
+    
+    User Message: {message}
+    
+    Please respond considering both the message and any attached documents:
+    """
 
-    # Configure generation parameters
-    generation_params = dict(
-        temperature=temperature,
-        max_new_tokens=int(max_tokens),
-        top_p=top_p,
-        repetition_penalty=1.1,  # If you want to add a penalty
+    generate_kwargs = dict(
+        temperature=0.7,
+        max_new_tokens=2000,
+        top_p=0.95,
+        repetition_penalty=1.2,
     )
-
-    # Stream the response
-    partial_response = ""
+    
+    # Generate response
     stream = client.text_generation(
-        prompt,
+        full_prompt,
         stream=True,
         details=True,
-        **generation_params
+        **generate_kwargs
     )
-    for chunk in stream:
-        if chunk.token.special:
+    
+    partial_message = ""
+    for response in stream:
+        if response.token.special:
             continue
-        partial_response += chunk.token.text
-        # Yield partial chunks so user sees a "streaming" effect
-        yield partial_response
-
-
-###############################################################################
-# BUILD THE GRADIO INTERFACE
-###############################################################################
+        partial_message += response.token.text
+        yield partial_message
 
 with gr.Blocks(theme="soft") as demo:
-    # Top Title or Branding
-    gr.Markdown(
-        """
-        <h1 style="text-align:center; margin-bottom: 5px;">
-            <b>balianone.com</b>
-        </h1>
-        <h3 style="text-align:center; margin-top: 0px;">
-            Chat with DeepSeek-R1
-        </h3>
-        """,
-        elem_id="title"
-    )
+    gr.Markdown("# DeepSeek-R1 Assistant with File Support")
+    gr.Markdown("Upload images, PDFs, or docs and chat about them!")
     
     with gr.Row():
-        # Left Column (System message, HF API key, sliders)
-        with gr.Column(scale=1, min_width=250):
-            system_message = gr.Textbox(
-                label="System message", 
-                value="You are a friendly Chatbot created by Bali"
-            )
-            user_api_key = gr.Textbox(
-                label="HF API Key (optional)", 
-                type="password",
-                placeholder="Leave blank to use open-source DeepSeek key"
-            )
-            
-            gr.Markdown(
-                "If error occurs, you can set your own API key. You’ll get R1 premium access."
-            )
-            
-            max_tokens = gr.Slider(
-                minimum=1, 
-                maximum=4000, 
-                step=1, 
-                value=4000, 
-                label="Max new tokens"
-            )
-            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.01,
-                label="Top-p (nucleus sampling)"
-            )
-
-        # Right Column (Chat interface)
-        with gr.Column(scale=3):
-            chatbot = gr.ChatInterface(
-                fn=chat_with_deepseek, 
-                # Additional inputs that feed into your chat function:
-                additional_inputs=[system_message, user_api_key, max_tokens, temperature, top_p],
-                
-                # Use the newer messages format to avoid warnings
-                type="messages",
-                
-                # You can customize how the chat area looks here as well:
-                chatbot_base_url=None,  # Keep local
-                height=550,
-                title="DeepSeek-R1 Assistant"
-            )
-
-    # Launch the app
-    demo.launch()
-
-
+        file_input = gr.File(label="Upload File (PDF/Image/Doc)", type="filepath")
+        
+    chatbot = gr.ChatInterface(
+        fn=chat,
+        additional_inputs=[file_input],
+        type="messages",
+        examples=[]
+    )
 
+demo.launch()