Update app.py

app.py

@@ -1,42 +1,35 @@
 import gradio as gr
 from datasets import load_dataset
-from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 from sklearn.feature_extraction.text import TfidfVectorizer
 from sklearn.metrics.pairwise import cosine_similarity
 import numpy as np
 
 # Load dataset
 ds = load_dataset("STEM-AI-mtl/Electrical-engineering", split="train")
-docs = [q + " " + a for q, a in zip(ds["input"], ds["output"])]
+
+# Combine input (question) and output (answer) for vectorization
+docs = [inp + " " + out for inp, out in zip(ds["input"], ds["output"])]
 vectorizer = TfidfVectorizer().fit(docs)
 tfidf_matrix = vectorizer.transform(docs)
 
-# Load LLM
-model_name = "STEM-AI-mtl/phi-2-electrical-engineering"
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(model_name)
-qa_pipeline = pipeline("text-generation", model=model, tokenizer=tokenizer, max_length=512)
-
 # Retrieval function
-def retrieve(user_q, top_k=1):
+def retrieve_answer(user_q, top_k=1):
     vec = vectorizer.transform([user_q])
     sims = cosine_similarity(vec, tfidf_matrix).flatten()
-    best = np.argmax(sims)
-    return ds["input"][best], ds["output"][best]
+    idxs = np.argsort(-sims)[:top_k]
+    return "\n\n".join([f"**Q:** {ds['input'][i]}\n**A:** {ds['output'][i]}" for i in idxs])
 
-# Full QA function
-def answer(user_q):
-    q, a = retrieve(user_q)
-    prompt = f"Question: {user_q}\nRelevant Q: {q}\nRelevant A: {a}\nAnswer:"
-    gen = qa_pipeline(prompt)
-    return gen[0]["generated_text"].split("Answer:")[-1].strip()
-
-# Gradio UI
+# Gradio app
 with gr.Blocks() as demo:
-    gr.Markdown("# 🛠️ EE Smart Q&A Chatbot")
-    user_q = gr.Textbox(label="Ask your electronics question", lines=2)
-    out = gr.Markdown()
-    user_q.submit(answer, inputs=user_q, outputs=out)
+    gr.Markdown("# 🤖 Electronics Engineering Q&A Chatbot")
+    gr.Markdown("Ask any electronics-related question and get an AI-assisted answer based on a curated dataset.")
+    
+    with gr.Row():
+        user_q = gr.Textbox(label="Your Question", lines=2, placeholder="e.g. What is the purpose of a Zener diode?")
+    answer_box = gr.Markdown(label="Answer")
+    submit_btn = gr.Button("Get Answer")
+
+    submit_btn.click(fn=retrieve_answer, inputs=user_q, outputs=answer_box)
 
 if __name__ == "__main__":
     demo.launch()