Update app.py

app.py

@@ -1,5 +1,5 @@
 import streamlit as st
-from langchain.chat_models import HuggingFaceHub
+from transformers import pipeline
 from langchain.schema import AIMessage, HumanMessage, SystemMessage
 
 # ------------------------
@@ -14,21 +14,24 @@ if "sessionMessages" not in st.session_state:
     ]
 
 # ------------------------
-# Load Hugging Face Model
+# Load Hugging Face pipeline
 # ------------------------
-chat = HuggingFaceHub(
-    repo_id="OpenAssistant/oasst-sft-4-pythia-12b",
-    model_kwargs={"temperature": 0, "max_new_tokens": 512}
+# Using a free, instruction-tuned model
+generator = pipeline(
+    "text-generation",
+    model="OpenAssistant/oasst-sft-4-pythia-12b",
+    device=0,  # GPU if available
+    max_new_tokens=512,
+    temperature=0
 )
 
 # ------------------------
 # Helper Functions
 # ------------------------
 def load_answer(question):
-    # Add human message
     st.session_state.sessionMessages.append(HumanMessage(content=question))
 
-    # Convert messages to a single prompt string
+    # Convert session messages to a single string prompt
     prompt = ""
     for msg in st.session_state.sessionMessages:
         if isinstance(msg, SystemMessage):
@@ -38,13 +41,13 @@ def load_answer(question):
         elif isinstance(msg, AIMessage):
             prompt += f"AI: {msg.content}\n"
 
-    # Get AI response
-    assistant_answer = chat(prompt)
+    # Generate response
+    output = generator(prompt, max_new_tokens=512, do_sample=True, temperature=0)
+    answer_text = output[0]["generated_text"][len(prompt):].strip()  # remove prompt from output
 
-    # Add AI message to session
-    st.session_state.sessionMessages.append(AIMessage(content=assistant_answer.content))
+    st.session_state.sessionMessages.append(AIMessage(content=answer_text))
 
-    return assistant_answer.content
+    return answer_text
 
 def get_text():
     return st.text_input("You: ", key="input")