Update app.py

app.py

@@ -1,31 +1,33 @@
-import os
 import gradio as gr
 import pandas as pd
 import re
-from langchain_community.vectorstores import FAISS
 from langchain_community.embeddings import HuggingFaceEmbeddings
-from huggingface_hub import InferenceClient
+from langchain_community.vectorstores import FAISS
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import torch
 
-# === Paths ===
+# --- Paths ---
 CSV_FOLDER = "data"
 FAISS_INDEX_PATH = "data/faiss_index_hugface_BAAI_new"
 
-# === Load CSVs ===
+# --- Load CSVs ---
 d1 = pd.read_csv(f"{CSV_FOLDER}/dataset1_clean.csv")
 d2 = pd.read_csv(f"{CSV_FOLDER}/dataset2_clean.csv")
 d3 = pd.read_csv(f"{CSV_FOLDER}/dataset3_clean.csv")
 print("✅ CSVs loaded")
 
-# === Load FAISS ===
+# --- Load FAISS with dummy embeddings ---
 embeddings = HuggingFaceEmbeddings(model_name="BAAI/bge-base-en")
 faiss_index = FAISS.load_local(FAISS_INDEX_PATH, embeddings, allow_dangerous_deserialization=True)
 print("✅ FAISS loaded")
 
-# === Hugging Face Inference API ===
-HF_API_TOKEN = os.environ.get("HF_API_TOKEN")  # set this in Space secrets or locally
-client = InferenceClient(model="mistralai/Mistral-7B-Instruct-v0.3", token=HF_API_TOKEN)
+# --- Load Mistral model ---
+MODEL_ID = "mistralai/Mistral-7B-Instruct-v0.3"
+tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
+model = AutoModelForCausalLM.from_pretrained(MODEL_ID, torch_dtype=torch.bfloat16, device_map="auto")
+print("✅ Mistral model loaded")
 
-# === Property synonyms ===
+# --- Property synonyms ---
 property_synonyms = {
     "hardness": ["hv", "hardness", "vicker's hardness", "vickers hardness"],
     "bulk modulus": ["d_bulk (gpa)", "bulk modulus", "bulk_modulus"],
@@ -35,6 +37,7 @@ property_synonyms = {
     "density": ["density_exp", "density_calc", "density"]
 }
 
+# --- Helper functions ---
 def find_column_for_property(df, property_name):
     synonyms = property_synonyms.get(property_name.lower(), [property_name])
     for syn in synonyms:
@@ -86,8 +89,8 @@ def filter_all_datasets(datasets, queries, top_n=10):
     for df, name in datasets:
         df_filtered = df.copy()
         phase_filter = queries.get("phase", None)
-        phase_col = None
         if phase_filter:
+            phase_col = None
             for col in df_filtered.columns:
                 if any(phase_key in col.lower() for phase_key in property_synonyms["phase"]):
                     phase_col = col
@@ -107,26 +110,29 @@ def filter_all_datasets(datasets, queries, top_n=10):
             df_filtered = apply_numeric_filter(df_filtered, col, filter_val)
         if df_filtered is None or df_filtered.empty:
             continue
-        show_cols = ["formula"] if "formula" in df_filtered.columns else []
+        show_cols = []
+        if "formula" in df_filtered.columns:
+            show_cols.append("formula")
         for prop in queries.keys():
             if prop == "phase":
                 continue
             col = find_column_for_property(df_filtered, prop)
-            if col and col in df_filtered.columns and col not in show_cols:
+            if col and col in df_filtered.columns:
                 show_cols.append(col)
-        if phase_filter and phase_col and phase_col not in show_cols:
+        if phase_filter and phase_col:
             show_cols.append(phase_col)
+        show_cols = [c for c in show_cols if c in df_filtered.columns]
         df_filtered = df_filtered[show_cols].head(top_n).copy()
         df_filtered["Source"] = name
         results[name] = df_filtered
     return results
 
-# === Main Query Function ===
+# --- Main HEA query function ---
 def query_hea(question, top_k=5):
     # FAISS retrieval
     faiss_results = faiss_index.similarity_search(question, k=top_k)
     faiss_text = "\n".join([doc.page_content for doc in faiss_results])
-    
+
     # CSV filtering
     queries = parse_query_to_filters(question)
     csv_results_dict = filter_all_datasets(
@@ -134,12 +140,11 @@ def query_hea(question, top_k=5):
         queries,
         top_n=top_k
     )
-    
     csv_context = ""
     for name, df_filtered in csv_results_dict.items():
         csv_context += f"\n### {name} matches:\n{df_filtered.to_string(index=False)}\n"
 
-    # Prompt for Mistral Instruct conversational API
+    # --- Prompt for Mistral ---
     prompt = f"""
 You are a materials scientist. Based on the following context, answer precisely.
 FAISS context: {faiss_text}
@@ -147,18 +152,22 @@ CSV datasets context: {csv_context}
 Question: {question}
 Answer:
 """
-    # Conversational API requires role-based input
-    conversation_input = [{"role": "user", "content": prompt}]
-    response = client.conversation(conversation_input)
-    output_text = response[0]["generated_text"]
 
+    # Tokenize and generate
+    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+    outputs = model.generate(**inputs, max_new_tokens=512, do_sample=True, temperature=0.0)
+    answer = tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+    # Merge CSV results
     merged_df = pd.concat(csv_results_dict.values(), ignore_index=True) if csv_results_dict else pd.DataFrame()
-    return output_text, merged_df, faiss_text
 
+    return answer, merged_df, faiss_text
+
+# --- Gradio wrapper ---
 def gradio_query(question):
     return query_hea(question)
 
-# === Gradio Interface ===
+# --- Launch Gradio interface ---
 demo = gr.Interface(
     fn=gradio_query,
     inputs=gr.Textbox(lines=2, placeholder="Ask about HEAs..."),
@@ -171,6 +180,5 @@ demo = gr.Interface(
     description="Query HEA datasets + FAISS paper embeddings"
 )
 
-if __name__ == "__main__":
-    demo.launch()
+demo.launch()