Modified grok model

app.py

@@ -3,7 +3,7 @@ from graph import build_graph
 from utils import get_retriever, load_vectorstore_from_text
 from pypdf import PdfReader
 import hashlib
-from transformers import pipeline
+from transformers import pipeline, BartForConditionalGeneration, BartTokenizer
 
 # --- Page Config ---
 st.set_page_config(page_title="LangGraph RAG Chatbot", layout="wide")
@@ -82,9 +82,12 @@ with st.sidebar:
 
 # --- Initialize Summarizer ---
 if "summarizer" not in st.session_state:
+    tokenizer = BartTokenizer.from_pretrained("facebook/bart-large-cnn")
+    model = BartForConditionalGeneration.from_pretrained("facebook/bart-large-cnn")
     st.session_state.summarizer = pipeline(
         "summarization",
-        model="facebook/bart-large-cnn",
+        model=model,
+        tokenizer=tokenizer,
         device=-1
     )
 
@@ -106,50 +109,48 @@ if "history" not in st.session_state:
     st.session_state.history = []
 
 # --- Query Input ---
-if "current_query" not in st.session_state:
-    st.session_state.current_query = ""
+query_input = st.text_input("💬 Ask a question:")
 
-query = st.text_input("💬 Ask a question:", key="current_query")
 send_triggered = st.button("Send")
 
-# --- Send Query ---
-if send_triggered and query.strip():
-    formatted_history = [(q, r) for q, r, _ in st.session_state.history]
-
-    with st.spinner("Generating response..."):
-        try:
-            result = st.session_state.graph(
-                query=query,
-                temperature=temperature,
-                raw_text=st.session_state.get("raw_text"),
-                history=formatted_history,
-                retriever_override=st.session_state.get("retriever")
-            )
-
-            response = result.get("response", "No response generated.")
-            retrieved_docs = result.get("retrieved_docs", [])
-
-            st.markdown("### 🤖 Response")
-            st.markdown(response)
-
-            # Save to history
-            st.session_state.history.append((query, response, retrieved_docs))
-
-            # Show retrieved docs
-            if retrieved_docs:
-                with st.expander("📄 Retrieved Chunks"):
-                    for j, doc in enumerate(retrieved_docs):
-                        content = getattr(doc, "text", str(doc))
-                        st.markdown(f"**Chunk {j+1}:**")
-                        st.code(content.strip(), language="markdown")
-
-            # Clear input
-            st.session_state.current_query = ""
-
-        except Exception as e:
-            st.error(f"Query failed: {e}")
-elif send_triggered:
-    st.warning("Please enter a question.")
+if send_triggered:
+    if query_input.strip():
+        formatted_history = [(q, r) for q, r, _ in st.session_state.history]
+
+        with st.spinner("Generating response..."):
+            try:
+                result = st.session_state.graph(
+                    query=query_input,
+                    temperature=temperature,
+                    raw_text=st.session_state.get("raw_text"),
+                    history=formatted_history,
+                    retriever_override=st.session_state.get("retriever")
+                )
+
+                response = result.get("response", "No response generated.")
+                retrieved_docs = result.get("retrieved_docs", [])
+
+                st.markdown("### 🤖 Response")
+                st.markdown(response)
+
+                # Save to history
+                st.session_state.history.append((query_input, response, retrieved_docs))
+
+                # Show retrieved docs if available
+                if retrieved_docs:
+                    with st.expander("📄 Retrieved Chunks"):
+                        for j, doc in enumerate(retrieved_docs):
+                            content = getattr(doc, "text", str(doc))
+                            st.markdown(f"**Chunk {j+1}:**")
+                            st.code(content.strip(), language="markdown")
+
+                # Clear the input field by rerunning widget with empty value
+                st.experimental_rerun()
+
+            except Exception as e:
+                st.error(f"Query failed: {e}")
+    else:
+        st.warning("Please enter a question.")
 
 # --- Chat History Display ---
 if st.session_state.history:

graph.py

@@ -9,7 +9,7 @@ from langgraph.graph import StateGraph, END
 from llama_index.core import VectorStoreIndex
 from llama_index.core.retrievers import BaseRetriever
 from langchain_groq import ChatGroq
-from transformers import pipeline as hf_pipeline
+from transformers import pipeline, BartForConditionalGeneration, BartTokenizer
 
 # --- 1. Define the State for the Graph ---
 class GraphState(TypedDict):
@@ -180,7 +180,7 @@ def build_graph(model_type: str = "groq", retriever=None, summarizer=None):
         if not api_key:
             raise ValueError("GROQ_API_KEY environment variable not set.")
         llm = ChatGroq(
-            model="mixtral-8x7b-32768",  
+            model="x-ai/grok-4-fast:free",  
             api_key=api_key,
             temperature=0.7,
         )
@@ -195,10 +195,15 @@ def build_graph(model_type: str = "groq", retriever=None, summarizer=None):
         )
     else:
         raise ValueError("Invalid model_type. Choose 'groq' or 'gemini'.")
+    
+    def get_default_summarizer():
+        tokenizer = BartTokenizer.from_pretrained("facebook/bart-large-cnn")
+        model = BartForConditionalGeneration.from_pretrained("facebook/bart-large-cnn")
+        return pipeline("summarization", model=model, tokenizer=tokenizer, device=-1)
 
     if summarizer is None:
         print("---NO SUMMARIZER PROVIDED, USING DEFAULT (facebook/bart-large-cnn)---")
-        summarizer = hf_pipeline("summarization", model="facebook/bart-large-cnn")
+        summarizer = get_default_summarizer()
 
     workflow = StateGraph(GraphState)