fix

app.py

@@ -5,36 +5,32 @@ from datasets import load_dataset
 from sentence_transformers import SentenceTransformer
 from transformers import AutoTokenizer, AutoModelForCausalLM
 
+dataset = load_dataset("ariG23498/pis-blogs-chunked")
+embedding_model = SentenceTransformer(model_name_or_path="all-mpnet-base-v2")
+tokenizer = AutoTokenizer.from_pretrained("google/gemma-2b-it")
+model = AutoModelForCausalLM.from_pretrained("google/gemma-2b-it", torch_dtype=torch.bfloat16)
+
 @spaces.GPU
 def process_query(query):
-    dataset = load_dataset("ariG23498/pis-blogs-chunked")
-    embedding_model = SentenceTransformer(model_name_or_path="all-mpnet-base-v2", device="cuda:0")
+    embedding_model = embedding_model.to("cuda")
     text_embeddings = embedding_model.encode(dataset["train"]["text"])
-    tokenizer = AutoTokenizer.from_pretrained("google/gemma-2b-it")
-    model = AutoModelForCausalLM.from_pretrained("google/gemma-2b-it", torch_dtype=torch.bfloat16, device_map="auto")
 
-    print(f"Query: {query}")
-    
     query_embedding = embedding_model.encode(query)
-    
     similarity_scores = embedding_model.similarity(query_embedding, text_embeddings)
     top_indices = (-similarity_scores).argsort()[0][:5]
     
     context = dataset["train"]["text"][top_indices[0]]
     url = dataset["train"]["url"][top_indices[0]]
     
-    print(f"Searching URL: {url}")
-    print(f"Found context: {context}")
-    
     input_text = (
         f"Based on the context provided, '{context}', how would"
         f"you address the user's query regarding '{query}'? Please"
         " provide a detailed and contextually relevant response."
     )
     
-    input_ids = tokenizer(input_text, return_tensors="pt").to("cuda:0")
+    input_ids = tokenizer(input_text, return_tensors="pt").to("cuda")
     len_text = len(input_text)
-    
+    model = model.to("cuda")
     with torch.inference_mode():
         generated_outputs = model.generate(**input_ids, max_new_tokens=1000, do_sample=False)
         generated_outputs = tokenizer.batch_decode(generated_outputs, skip_special_tokens=True)