Update app.py

app.py

@@ -1,77 +1,36 @@
-import streamlit as st
-from langchain.prompts import ChatPromptTemplate  # type: ignore
-from langchain_huggingface import HuggingFacePipeline  # type: ignore
-import json
-import os
+# Import necessary modules
+from langchain.prompts import ChatPromptTemplate # type: ignore
+from langchain.llms import Ollama # type: ignore
+import streamlit as st # type: ignore
 
 # Streamlit setup
 st.title("Subbu Chat Bot")
-
-# Initialize session state for storing conversation history
-if "history" not in st.session_state:
-    st.session_state.history = []
-
-# User input
-input_txt = st.text_input("Please enter your queries here...")
+input_txt = st.text_input("Enter your queries here...")
 
 # Add a dropdown for model selection
-model_choice = st.selectbox("Select the model:", ["Meta-Llama-8B"])
+model_choice = st.selectbox("Select the model:", ["Llama 3.2", "Llama 3.1", "Code Llama", "subbu"])
 
 # Define the prompt template
-prompt_template = ChatPromptTemplate.from_messages(
-    [("system", "You are a helpful AI assistant. Your name is Subbu Assistant.")]
-)
-
-# Initialize the model using Hugging Face Transformers
-llm_8b = HuggingFacePipeline.from_model_id(
-    model_id="facebook-llama/Meta-Llama-8B",
-    task="text-generation",
-    pipeline_kwargs={"max_new_tokens": 100, "top_k": 50, "temperature": 0.7},
-    api_token="your_hugging_face_api_token"
+prompt = ChatPromptTemplate.from_messages(
+    [("system", "You are a helpful AI assistant. Your name is Subbu Assistant."),
+     ("user", "user query: {query}")]
 )
 
-# Function to get the response from the assistant
-def get_response(model, input_text):
-    prompt = prompt_template.format(query=input_text)
-    response = model(prompt)
-    return response[0]['generated_text']
-
-# Function to save session state to a file
-def save_session_state():
-    with open("session_state.json", "w") as f:
-        json.dump(st.session_state.history, f)
-
-# Function to load session state from a file
-def load_session_state():
-    if os.path.exists("session_state.json"):
-        with open("session_state.json", "r") as f:
-            st.session_state.history = json.load(f)
-
-# Load session state if it exists
-load_session_state()
+# Initialize each model (adjust the model names based on available models)
+llm_3_2 = Ollama(model="llama3.2")
+llm_3_1 = Ollama(model="llama3.1")
+code_llama = Ollama(model="codellama")
+subbu =  Ollama(model="subbu")
 
 # Process input and display the response
-if st.button("Send"):
-    if input_txt:
-        # Select model based on user choice
-        model = llm_8b
-
-        # Get the model response
-        response = get_response(model, input_txt)
-
-        # Save the query and response in the session history
-        st.session_state.history.append({"user": input_txt, "assistant": response})
-
-        # Save session state
-        save_session_state()
-
-# Clear conversation button
-if st.button("Clear Conversation"):
-    st.session_state.history = []
-    if os.path.exists("session_state.json"):
-        os.remove("session_state.json")
-
-# Display conversation history
-for entry in st.session_state.history:
-    st.write(f"**User:** {entry['user']}")
-    st.write(f"**Subbu Assistant:** {entry['assistant']}")
+if input_txt:
+    # Select model based on user choice
+    if model_choice == "Llama 3.2":
+        response = llm_3_2(prompt.format(query=input_txt))
+    elif model_choice == "Llama 3.1":
+        response = llm_3_1(prompt.format(query=input_txt))
+    elif model_choice == "Code Llama":
+        response = code_llama(prompt.format(query=input_txt))
+    
+    # Display the response
+    st.write(response)