Update app.py

app.py

@@ -1,6 +1,5 @@
 import gradio as gr
 import torch
-import spaces
 from langchain_huggingface import HuggingFaceEmbeddings
 from langchain_community.document_loaders import TextLoader
 from langchain_community.vectorstores import FAISS
@@ -8,6 +7,10 @@ from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain.chains import RetrievalQA
 from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 from langchain_huggingface import HuggingFacePipeline
+import spaces
+zero = torch.Tensor([0]).cuda()
+print(zero.device)  # This will likely print 'cpu'
+device = "cuda" if torch.cuda.is_available() else "cpu"
 
 # Load and process the document
 doc_loader = TextLoader("dataset.txt")
@@ -22,18 +25,11 @@ vectordb = FAISS.from_documents(split_docs, embeddings)
 # Load model and tokenizer
 model_name = "01-ai/Yi-Coder-9B-Chat"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
-
-@spaces.GPU
-def setup_model():
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    model = AutoModelForCausalLM.from_pretrained(
-        model_name, 
-        device_map="auto", 
-        torch_dtype=torch.float16 if device.type == "cuda" else torch.float32
-    )
-    return model, device
-
-model, device = setup_model()
+model = AutoModelForCausalLM.from_pretrained(
+    model_name, 
+    device_map="auto", 
+    torch_dtype=torch.float16 if device == "cuda" else torch.float32
+)
 
 # Set up the QA pipeline
 qa_pipeline = pipeline(
@@ -41,8 +37,7 @@ qa_pipeline = pipeline(
     model=model,
     tokenizer=tokenizer,
     max_new_tokens=750,
-    pad_token_id=tokenizer.eos_token_id,
-    device=device
+    pad_token_id=tokenizer.eos_token_id
 )
 
 llm = HuggingFacePipeline(pipeline=qa_pipeline)
@@ -67,7 +62,6 @@ def clean_response(response):
         return result.split("Answer:")[1].strip()
     return result.strip()
 
-@spaces.GPU
 def chatbot_response(user_input):
     processed_query = preprocess_query(user_input)
     raw_response = qa_chain.invoke({"query": processed_query})
@@ -90,4 +84,4 @@ with gr.Blocks() as chat_interface:
 
 # Launch the interface
 if __name__ == "__main__":
-    chat_interface.launch()
+    chat_interface.launch()