Update app.py

app.py

@@ -1,9 +1,9 @@
 import spaces
 import gradio as gr
 import torch
-from langchain_huggingface import HuggingFaceEmbeddings, HuggingFacePipeline
 from langchain_community.vectorstores import FAISS
-from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
+from langchain_huggingface import HuggingFaceEmbeddings, HuggingFacePipeline
+from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer
 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
@@ -14,7 +14,6 @@ from langchain_core.prompts import PromptTemplate
 print("Loading FAISS Vector Store...")
 embeddings = HuggingFaceEmbeddings(model_name='sentence-transformers/all-MiniLM-L6-v2')
 
-# allow_dangerous_deserialization is required for FAISS in newer LangChain versions
 vectorstore = FAISS.load_local(
     "faiss_upf_index", 
     embeddings, 
@@ -23,21 +22,16 @@ vectorstore = FAISS.load_local(
 retriever = vectorstore.as_retriever()
 
 # ==========================================
-# 2. Load Model & Tokenizer in 4-bit
+# 2. Load Model & Tokenizer in 16-bit (ZeroGPU Fix)
 # ==========================================
 print("Loading Model...")
 model_id = "anirudh248/llama3-upf-generator"
 
-# Load in 4-bit to fit perfectly inside a GPU Space
-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,
-    quantization_config=bnb_config,
+    torch_dtype=torch.bfloat16,
     device_map="auto"
 )
 model.generation_config.pad_token_id = tokenizer.eos_token_id
@@ -74,16 +68,15 @@ prompt = PromptTemplate.from_template(prompt_template)
 document_chain = create_stuff_documents_chain(llm, prompt)
 rag_chain = create_retrieval_chain(retriever, document_chain)
 
+def generate_upf_code(power_intent_description):
+    result = rag_chain.invoke({"input": power_intent_description})
+    return result['answer'].strip()
 
 # ==========================================
-# 4. Gradio UI & Inference
+# 4. Gradio UI
 # ==========================================
 
 @spaces.GPU(duration=120) 
-def generate_upf_code(power_intent_description):
-    result = rag_chain.invoke({"input": power_intent_description})
-    return result['answer'].strip()
-
 def user_interaction(user_message, history):
     history = history or []
     response = generate_upf_code(user_message)