Update app.py

app.py

@@ -3,15 +3,15 @@ from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 import PyPDF2
 import torch
 
-st.set_page_config(page_title="Perplexity Clone (Gemma)", layout="wide")
-st.title("📚 Perplexity-Style AI Study Assistant using Gemma")
+st.set_page_config(page_title="Perplexity-style Q&A (Mistral)", layout="wide")
+st.title("🧠 Perplexity-style AI Study Assistant using Mistral 7B")
 
-# Load Gemma model and tokenizer
+# Load Mistral model and tokenizer
 @st.cache_resource
 def load_model():
-    tokenizer = AutoTokenizer.from_pretrained("google/gemma-7b-it")
+    tokenizer = AutoTokenizer.from_pretrained("mistralai/Mistral-7B-Instruct-v0.1")
     model = AutoModelForCausalLM.from_pretrained(
-        "google/gemma-7b-it",
+        "mistralai/Mistral-7B-Instruct-v0.1",
         torch_dtype=torch.float16,
         device_map="auto"
     )
@@ -29,20 +29,22 @@ def extract_text_from_pdf(file):
     return text.strip()
 
 # UI Layout
-query = st.text_input("Ask a question or type a query:")
+query = st.text_input("Ask a question or enter a topic:")
 
-uploaded_file = st.file_uploader("Or upload a PDF to analyze its content:", type=["pdf"])
+uploaded_file = st.file_uploader("Or upload a PDF to use as context:", type=["pdf"])
 
 context = ""
 if uploaded_file:
     context = extract_text_from_pdf(uploaded_file)
-    st.text_area("Extracted Content", context, height=200)
+    st.text_area("📄 Extracted PDF Text", context, height=200)
 
 if st.button("Generate Answer"):
-    with st.spinner("Generating with Gemma..."):
+    with st.spinner("Generating answer with Mistral 7B..."):
         prompt = query
         if context:
-            prompt = f"Context:\n{context}\n\nQuestion: {query}"
+            prompt = f"[INST] Use the following context to answer the question:\n\n{context}\n\nQuestion: {query} [/INST]"
+        else:
+            prompt = f"[INST] {query} [/INST]"
         output = textgen(prompt)[0]["generated_text"]
         st.success("Answer:")
-        st.write(output)
+        st.write(output.replace(prompt, "").strip())