Create app.py

app.py

@@ -0,0 +1,89 @@
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import torch
+
+# Load DistilGPT-2 (lightweight and fast)
+distilgpt2_tokenizer = AutoTokenizer.from_pretrained("distilgpt2")
+distilgpt2_model = AutoModelForCausalLM.from_pretrained("distilgpt2")
+
+# Move model to GPU if available
+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)
+    text = ""
+    for page in doc:
+        text += page.get_text()
+    return text
+
+# Function to extract text from image
+def extract_text_from_image(image_path):
+    reader = easyocr.Reader(['en'])
+    results = reader.readtext(image_path)
+    extracted_text = " ".join([res[1] for res in results])
+    return extracted_text
+
+def chatbot(input_type, input_data):
+    if input_type == "Text":
+        prompt = input_data
+    elif input_type == "PDF":
+        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":
+        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
+    response = generate_response(prompt)
+    return response
+
+# List of input types
+INPUT_TYPES = ["Text", "PDF", "Image"]
+
+# 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")
+    ],
+    outputs="text",
+    title="Lightweight Chatbot with PDF and Image Support",
+    description="Select the input type (Text, PDF, or Image) and ask your question!"
+)
+
+# Launch the app
+interface.launch()