Update app.py

app.py

@@ -1,8 +1,8 @@
+import spaces
 import gradio as gr
 import torch
-from langchain_community.embeddings import HuggingFaceEmbeddings
+from langchain_huggingface import HuggingFaceEmbeddings, HuggingFacePipeline
 from langchain_community.vectorstores import FAISS
-from langchain_huggingface import HuggingFacePipeline
 from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
 from langchain_classic.chains import create_retrieval_chain
 from langchain_classic.chains.combine_documents import create_stuff_documents_chain
@@ -28,7 +28,7 @@ retriever = vectorstore.as_retriever()
 print("Loading Model...")
 model_id = "anirudh248/llama3-upf-generator"
 
-# Load in 4-bit to fit inside a T4 GPU Space
+# Load in 4-bit to fit perfectly inside a GPU Space
 bnb_config = BitsAndBytesConfig(
     load_in_4bit=True,
     bnb_4bit_compute_dtype=torch.float16
@@ -74,13 +74,16 @@ prompt = PromptTemplate.from_template(prompt_template)
 document_chain = create_stuff_documents_chain(llm, prompt)
 rag_chain = create_retrieval_chain(retriever, document_chain)
 
+
+# ==========================================
+# 4. Gradio UI & Inference
+# ==========================================
+
+@spaces.GPU(duration=120) 
 def generate_upf_code(power_intent_description):
     result = rag_chain.invoke({"input": power_intent_description})
     return result['answer'].strip()
 
-# ==========================================
-# 4. Gradio UI
-# ==========================================
 def user_interaction(user_message, history):
     history = history or []
     response = generate_upf_code(user_message)
@@ -95,7 +98,6 @@ def user_interaction(user_message, history):
 with gr.Blocks() as interface:
     gr.Markdown("# ⚡ UPF Code Generator with Llama 3 & RAG")
     
-    # Removed `type="messages"` parameter to fix TypeError
     chatbot = gr.Chatbot(label="Chat History", elem_id="chatbot")
     
     with gr.Row():