Update app.py

app.py

@@ -17,13 +17,13 @@ def load_model():
     tokenizer = AutoTokenizer.from_pretrained(model_name)
     # Check if GPU is available
     device = "cuda" if torch.cuda.is_available() else "cpu"
-    model = AutoModel.from_pretrained(model_name).to("cuda")
-    return tokenizer, model
+    model = AutoModel.from_pretrained(model_name).to(device)
+    return tokenizer, model, device
 
-def generate_embedding(text, tokenizer, model):
+def generate_embedding(text, tokenizer, model, device):
     """Generate embeddings for a given text."""
     inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=512)
-    inputs = {k: v.to("cuda") for k, v in inputs.items()}
+    inputs = {k: v.to(device) for k, v in inputs.items()}
     with torch.no_grad():
         outputs = model.encoder(**inputs)
     return outputs.last_hidden_state.mean(dim=1).squeeze().cpu().numpy()
@@ -53,14 +53,14 @@ def main():
     st.write("Find Python repositories to learn production-level coding practices.")
 
     # Load resources
-    tokenizer, model = load_model()
+    tokenizer, model, device = load_model()
     data = load_data()
 
     # Input user query
     user_query = st.text_input("Describe your project or learning goal:",
                                "I am working on a project to recommend music using pandas and numpy.")
     if user_query:
-        query_embedding = generate_embedding(user_query, tokenizer, model)
+        query_embedding = generate_embedding(user_query, tokenizer, model, device)
 
         # Compute similarity
         data['similarity'] = data['embedding'].apply(