Update app.py

app.py

@@ -4,6 +4,7 @@ import os
 import faiss
 from sentence_transformers import SentenceTransformer
 from transformers import AutoTokenizer, AutoModelForCausalLM
+import torch
 
 # Load the knowledge base
 with open("knowledge_base.json", "r") as file:
@@ -29,41 +30,23 @@ llm = AutoModelForCausalLM.from_pretrained(llama_model_name, token=API_TOKEN)
 # Generate advice using RAG
 def generate_advice(extracted_data):
     try:
-        # Ensure extracted_data is valid
-        if not isinstance(extracted_data, list):
-            raise ValueError("Input data must be a list of dictionaries.")
-        if not all(isinstance(item, dict) for item in extracted_data):
-            raise ValueError("Each item in input data must be a dictionary.")
-
         recommendations = []
 
         for item in extracted_data:
-            # Validate required keys
+            # Validate input keys
             if not all(k in item for k in ["Component", "Status"]):
                 raise ValueError("Each input item must have 'Component' and 'Status' keys.")
-            
+
             # Prepare the query string
             query = f"{item['Component']} {item['Status']}"
-            print(f"Processing query: {query}")  # Debug print
+            print(f"Processing query: {query}")
 
-            # Generate query embedding and reshape
+            # Generate query embedding
             query_embedding = embedding_model.encode([query])
             query_embedding = np.array(query_embedding, dtype="float32").reshape(1, -1)
 
-            # Debugging embedding dimensions
-            print(f"Query Embedding Shape: {query_embedding.shape}, FAISS Index Dim: {index.d}")
-
-            # Validate embedding dimensions
-            if query_embedding.shape[1] != index.d:
-                raise ValueError(
-                    f"Embedding dimension mismatch: Query ({query_embedding.shape[1]}), Index ({index.d})"
-                )
-
             # Search for the closest match in FAISS
             _, idx = index.search(query_embedding, 1)
-            print(f"FAISS Index: {idx}, Best Match Raw: {kb[idx[0][0]]}")
-
-            # Retrieve the closest match
             best_match = kb[idx[0][0]]
 
             # Prepare the LLM prompt
@@ -78,21 +61,26 @@ def generate_advice(extracted_data):
             Provide additional insights or recommendations.
             """
 
-            # Generate advice using LLaMA model
+            # Tokenize input properly for LLaMA
             message_yours = [
                 {"role": "system", "content": role},
                 {"role": "user", "content": prompt},
             ]
 
+            # Properly tokenize to return a PyTorch tensor
             input_text_with_your_role = tokenizer.apply_chat_template(
                 message_yours,
-                tokenize=False,
+                tokenize=True,  # Must tokenize to return input_ids
                 add_generation_prompt=True,
                 return_tensors="pt",
             )
 
+            # Move tensor to appropriate device (CPU/GPU)
+            input_text_with_your_role = input_text_with_your_role.to(torch.device("cpu"))
+
+            # Generate advice
             output = llm.generate(
-                input_ids=input_text_with_your_role,
+                input_ids=input_text_with_your_role["input_ids"],
                 max_length=150,
                 num_return_sequences=1
             )
@@ -107,7 +95,7 @@ def generate_advice(extracted_data):
         return recommendations
 
     except Exception as e:
-        print(f"Error occurred: {str(e)}")  # Debugging error
+        print(f"Error occurred: {str(e)}")
         return [{"error": f"Exception occurred: {str(e)}"}]
 
 # Gradio app with LLM integration