Update app.py

app.py

@@ -1,6 +1,6 @@
 import torch
 import gradio as gr
-from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
+from transformers import AutoModelForCausalLM, AutoTokenizer
 from langchain.embeddings import HuggingFaceEmbeddings
 from langchain.vectorstores import FAISS
 from langchain.text_splitter import RecursiveCharacterTextSplitter
@@ -9,8 +9,9 @@ from langchain.document_loaders import PyPDFLoader
 # Model and Tokenizer
 MODEL_ID = "deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B"
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
 tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
-model = AutoModelForCausalLM.from_pretrained(MODEL_ID, device_map="auto")
+model = AutoModelForCausalLM.from_pretrained(MODEL_ID, torch_dtype=torch.float16, device_map="auto")
 
 # Load embedding model for RAG
 embedding_model = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
@@ -19,18 +20,24 @@ vector_store = None
 # Function to process PDF and create vector database
 def process_pdf(pdf_path):
     global vector_store
+    if not pdf_path:
+        return "❌ No PDF uploaded. Please upload a valid file."
+
     loader = PyPDFLoader(pdf_path)
     documents = loader.load()
     text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=100)
     texts = text_splitter.split_documents(documents)
     
+    if not texts:
+        return "❌ No text extracted from the PDF."
+
     vector_store = FAISS.from_documents(texts, embedding_model)
-    return "PDF successfully processed and indexed."
+    return "✅ PDF successfully processed and indexed."
 
 # RAG Query Function
-def query_rag(message, system_prompt, temperature, max_new_tokens, top_k, repetition_penalty, top_p, history=None):
+def query_rag(message, system_prompt, temperature, max_new_tokens, top_k, repetition_penalty, top_p, history=[]):
     if vector_store is None:
-        return "Please upload and process a PDF first."
+        return "⚠️ Please upload and process a PDF first."
     
     # Retrieve relevant chunks
     docs = vector_store.similarity_search(message, k=3)
@@ -53,7 +60,8 @@ def query_rag(message, system_prompt, temperature, max_new_tokens, top_k, repeti
         max_new_tokens=max_new_tokens,
         top_k=top_k,
         repetition_penalty=repetition_penalty,
-        top_p=top_p
+        top_p=top_p,
+        pad_token_id=tokenizer.eos_token_id  # Ensures correct padding
     )
     response = tokenizer.decode(output_ids[0], skip_special_tokens=True)
     return response
@@ -61,19 +69,23 @@ def query_rag(message, system_prompt, temperature, max_new_tokens, top_k, repeti
 # Gradio Interface
 def launch_interface():
     with gr.Blocks() as demo:
-        gr.Markdown("## 🤖 RAG Chatbot with DeepSeek")
-        pdf_uploader = gr.File(label="Upload PDF", type="filepath")
-        process_btn = gr.Button("Process PDF")
+        gr.Markdown("## 🤖 RAG Chatbot with DeepSeek AI")
+        
+        with gr.Row():
+            pdf_uploader = gr.File(label="📂 Upload PDF", type="filepath")
+            process_btn = gr.Button("📌 Process PDF")
+        
         process_output = gr.Textbox(label="Processing Status", interactive=False)
         
-        chatbot = gr.ChatInterface(query_rag,
+        chatbot = gr.ChatInterface(
+            fn=query_rag,
             additional_inputs=[
-                gr.Textbox("You are a helpful assistant.", label="System Prompt"),
-                gr.Slider(0, 1, 0.6, label="Temperature"),
-                gr.Slider(0, 32000, 10000, label="Max new tokens"),
-                gr.Slider(1, 80, 40, label="Top K"),
-                gr.Slider(0.1, 2.0, 1.1, label="Repetition Penalty"),
-                gr.Slider(0, 1, 0.95, label="Top P"),
+                gr.Textbox("You are a helpful assistant.", label="🔹 System Prompt"),
+                gr.Slider(0.1, 1.0, 0.6, label="🌡️ Temperature"),
+                gr.Slider(100, 32000, 1024, step=100, label="📏 Max new tokens"),
+                gr.Slider(1, 80, 40, step=1, label="🎯 Top K"),
+                gr.Slider(0.1, 2.0, 1.1, step=0.1, label="🔄 Repetition Penalty"),
+                gr.Slider(0.1, 1.0, 0.95, step=0.05, label="🔢 Top P"),
             ]
         )