Update app.py

app.py

@@ -1,19 +1,14 @@
-import spaces
 import gradio as gr
 import torch
 from langchain_community.vectorstores import FAISS
 from langchain_huggingface import HuggingFaceEmbeddings, HuggingFacePipeline
-from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer
+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
 from langchain_core.prompts import PromptTemplate
 
-# ==========================================
 # 1. Load Vector Store
-# ==========================================
-print("Loading FAISS Vector Store...")
 embeddings = HuggingFaceEmbeddings(model_name='sentence-transformers/all-MiniLM-L6-v2')
-
 vectorstore = FAISS.load_local(
     "faiss_upf_index", 
     embeddings, 
@@ -21,17 +16,19 @@ vectorstore = FAISS.load_local(
 )
 retriever = vectorstore.as_retriever()
 
-# ==========================================
-# 2. Load Model & Tokenizer in 16-bit
-# ==========================================
-print("Loading Model...")
+# 2. Load Model & Pipeline in 4-bit
 model_id = "anirudh248/llama3-upf-generator"
 
-tokenizer = AutoTokenizer.from_pretrained(model_id)
+bnb_config = BitsAndBytesConfig(
+    load_in_4bit=True,
+    bnb_4bit_compute_dtype=torch.float16
+)
 
+tokenizer = AutoTokenizer.from_pretrained(model_id)
 model = AutoModelForCausalLM.from_pretrained(
     model_id,
-    torch_dtype=torch.bfloat16
+    quantization_config=bnb_config,
+    device_map="auto"
 )
 model.generation_config.pad_token_id = tokenizer.eos_token_id
 
@@ -46,14 +43,9 @@ hf_pipeline = pipeline(
     return_full_text=False,
     clean_up_tokenization_spaces=False
 )
-
 llm = HuggingFacePipeline(pipeline=hf_pipeline)
 
-# ==========================================
-# 3. Setup Dual Chains
-# ==========================================
-
-# CHAIN A: The UPF Code Generator
+# 3. Setup RAG Chain & Chat Prompt
 upf_prompt_template = """
 You are an expert in Unified Power Format (UPF). Generate a precise and complete UPF code block based on the following power intent. Use the retrieved context as a reference. The code must be correct and adhere to UPF 3.0 standards.
 
@@ -69,11 +61,9 @@ upf_prompt = PromptTemplate.from_template(upf_prompt_template)
 document_chain = create_stuff_documents_chain(llm, upf_prompt)
 rag_chain = create_retrieval_chain(retriever, document_chain)
 
-# CHAIN B: The General Conversationalist
-chat_prompt = PromptTemplate.from_template("""You are a helpful AI assistant. Answer the user's message conversationally and concisely.
-
-User: {input}
-Assistant:""")
+chat_prompt = PromptTemplate.from_template(
+    "You are a helpful AI assistant. Answer the user's message conversationally and concisely.\n\nUser: {input}\nAssistant:"
+)
 
 def generate_upf_code(power_intent_description):
     result = rag_chain.invoke({"input": power_intent_description})
@@ -82,33 +72,24 @@ def generate_upf_code(power_intent_description):
 def generate_chat_response(user_message):
     return llm.invoke(chat_prompt.format(input=user_message)).strip()
 
-
-# ==========================================
-# 4. Gradio UI & Routing Logic
-# ==========================================
-
-# Keywords that trigger the high-quality UPF generation path
+# 4. Routing Logic & Gradio Interface
 UPF_KEYWORDS = [
     "upf", "power", "domain", "isolation", "retention", "voltage", 
     "switch", "supply", "pst", "level shifter", "state", "intent", 
     "create_", "set_", "connect_"
 ]
 
-@spaces.GPU(duration=120) 
 def user_interaction(user_message, history):
     history = history or []
-    
     msg_lower = user_message.lower()
     
     is_upf_request = any(kw in msg_lower for kw in UPF_KEYWORDS) or len(msg_lower.split()) > 12
 
     if is_upf_request:
-        # PATH A: High-quality UPF Generation
-        response = generate_upf_code(user_message)        
+        response = generate_upf_code(user_message)
         if "```" not in response:
             response = f"```tcl\n{response}\n```"
     else:
-        # PATH B: Normal Conversation
         response = generate_chat_response(user_message)
 
     history.append({"role": "user", "content": user_message})
@@ -116,7 +97,7 @@ def user_interaction(user_message, history):
     return history, ""
 
 with gr.Blocks() as interface:
-    gr.Markdown("# ⚡ UPF Code Generator with Llama 3 & RAG")
+    gr.Markdown("# UPF Code Generator with Llama 3 & RAG")
     
     chatbot = gr.Chatbot(label="Chat History", elem_id="chatbot")