Update app.py

app.py

@@ -1,7 +1,12 @@
-from transformers import AutoTokenizer, AutoModelForCausalLM
+import os
 import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import gradio as gr
+import fitz  # PyMuPDF
+import easyocr
+from PIL import Image
 
-# Load DistilGPT-2 (lightweight and fast)
+# Load a lightweight model (e.g., DistilGPT-2)
 distilgpt2_tokenizer = AutoTokenizer.from_pretrained("distilgpt2")
 distilgpt2_model = AutoModelForCausalLM.from_pretrained("distilgpt2")
 
@@ -9,36 +14,6 @@ distilgpt2_model = AutoModelForCausalLM.from_pretrained("distilgpt2")
 if torch.cuda.is_available():
     distilgpt2_model = distilgpt2_model.to("cuda")
 
-def generate_response(prompt):
-    # Tokenize the input prompt
-    inputs = distilgpt2_tokenizer(prompt, return_tensors="pt").to(distilgpt2_model.device)
-    
-    # Generate the response
-    outputs = distilgpt2_model.generate(**inputs, max_length=100)
-    
-    # Decode the response
-    response = distilgpt2_tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return response
-
-import gradio as gr
-
-# Gradio interface
-def chatbot(prompt):
-    response = generate_response(prompt)
-    return response
-
-interface = gr.Interface(
-    fn=chatbot,
-    inputs="text",
-    outputs="text",
-    title="DistilGPT-2 Chatbot",
-    description="Ask questions and get answers from DistilGPT-2!"
-)
-
-import fitz  # PyMuPDF
-import easyocr
-from PIL import Image
-
 # Function to extract text from PDF
 def extract_text_from_pdf(pdf_path):
     doc = fitz.open(pdf_path)
@@ -54,36 +29,70 @@ def extract_text_from_image(image_path):
     extracted_text = " ".join([res[1] for res in results])
     return extracted_text
 
+# Function to generate a response
+def generate_response(prompt):
+    inputs = distilgpt2_tokenizer(prompt, return_tensors="pt").to(distilgpt2_model.device)
+    outputs = distilgpt2_model.generate(**inputs, max_length=100)
+    response = distilgpt2_tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return response
+
+# Chatbot function to handle text, PDF, and image inputs
 def chatbot(input_type, input_data):
     if input_type == "Text":
         prompt = input_data
     elif input_type == "PDF":
+        if input_data is None:
+            return "Please upload a PDF file."
         pdf_text = extract_text_from_pdf(input_data)
         prompt = f"Extracted text from PDF:\n{pdf_text}\n\nQuestion: {input_data}"
     elif input_type == "Image":
+        if input_data is None:
+            return "Please upload an image file."
         image_text = extract_text_from_image(input_data)
         prompt = f"Extracted text from image:\n{image_text}\n\nQuestion: {input_data}"
     else:
         return "Invalid input type."
 
-    # Generate response using the selected model
+    # Generate response using the model
     response = generate_response(prompt)
     return response
 
-# List of input types
-INPUT_TYPES = ["Text", "PDF", "Image"]
-
 # Gradio interface
+input_components = [
+    gr.Dropdown(choices=["Text", "PDF", "Image"], label="Input Type"),
+    gr.Textbox(lines=2, placeholder="Enter text...", label="Text Input", visible=True),
+    gr.File(label="Upload PDF", file_types=[".pdf"], visible=False),
+    gr.Image(label="Upload Image", type="filepath", visible=False)
+]
+
+def update_input_components(input_type):
+    if input_type == "Text":
+        return [
+            gr.Textbox(visible=True),
+            gr.File(visible=False),
+            gr.Image(visible=False)
+        ]
+    elif input_type == "PDF":
+        return [
+            gr.Textbox(visible=False),
+            gr.File(visible=True),
+            gr.Image(visible=False)
+        ]
+    elif input_type == "Image":
+        return [
+            gr.Textbox(visible=False),
+            gr.File(visible=False),
+            gr.Image(visible=True)
+        ]
+
+# Create the Gradio interface
 interface = gr.Interface(
     fn=chatbot,
-    inputs=[
-        gr.Dropdown(choices=INPUT_TYPES, label="Input Type"),
-        gr.Textbox(lines=2, placeholder="Enter text or upload a file...", label="Input")
-    ],
+    inputs=input_components,
     outputs="text",
     title="Lightweight Chatbot with PDF and Image Support",
-    description="Select the input type (Text, PDF, or Image) and ask your question!"
+    description="Select the input type (Text, PDF, or Image) and provide your input."
 )
 
-# Launch the app
+# Add dynamic visibility for input components
 interface.launch()