Update app.py

app.py

@@ -1,41 +1,66 @@
-from langchain_openai import ChatOpenAI
-from langchain_core.prompts import ChatPromptTemplate
-from langchain_core.output_parsers import StrOutputParser
-from langchain_community.llms import Ollama
+# app.py
+
 import streamlit as st
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
 import os
 from dotenv import load_dotenv
 
 # Load environment variables
 load_dotenv()
 
-# Set environment variables
-os.environ["LANGCHAIN_TRACING_V2"] = "true"
-os.environ["LANGCHAIN_API_KEY"] = os.getenv("LANGCHAIN_API_KEY")
-
-# Prompt Template
-prompt = ChatPromptTemplate.from_messages(
-    [
-        ("system", "You are a helpful assistant. Please respond to the user queries"),
-        ("user", "Question: {question}")
-    ]
-)
+# Retrieve Hugging Face API token from environment variables (if accessing private models)
+HF_API_TOKEN = os.getenv("HF_API_TOKEN")  # Ensure you set this in Hugging Face Secrets
 
 # Streamlit app setup
-st.title('Langchain Demo With LLAMA2 API')
+st.title('Llama2 Chatbot Deployment on Hugging Face Spaces')
+st.write("This chatbot is powered by the Llama2 model. Ask me anything!")
+
+@st.cache_resource
+def load_model():
+    """
+    Load the tokenizer and model from Hugging Face.
+    This function is cached to prevent re-loading on every interaction.
+    """
+    tokenizer = AutoTokenizer.from_pretrained(
+        "meta-llama/Llama-2-7b-chat-hf",
+        use_auth_token=HF_API_TOKEN  # Remove if the model is public
+    )
+    model = AutoModelForCausalLM.from_pretrained(
+        "meta-llama/Llama-2-7b-chat-hf",
+        torch_dtype=torch.float16,  # Use float16 for reduced memory usage
+        device_map="auto",
+        use_auth_token=HF_API_TOKEN  # Remove if the model is public
+    )
+    return tokenizer, model
+
+# Load the model and tokenizer
+tokenizer, model = load_model()
 
 # User input
-input_text = st.text_input("Search the topic you want")
-
-# Ollama LLM (ensure the model is available, or access it through Hugging Face API)
-llm = Ollama(model="llama2")
-output_parser = StrOutputParser()
-chain = prompt | llm | output_parser
-
-# Display result when user inputs text
-if input_text:
-    try:
-        response = chain.invoke({"question": input_text})
-        st.write(response)
-    except Exception as e:
-        st.error(f"Error: {e}")
+user_input = st.text_input("You:", "")
+
+if user_input:
+    with st.spinner("Generating response..."):
+        try:
+            # Encode the input
+            inputs = tokenizer.encode(user_input + tokenizer.eos_token, return_tensors="pt").to(model.device)
+
+            # Generate a response
+            output = model.generate(
+                inputs,
+                max_length=1000,
+                temperature=0.7,
+                top_p=0.9,
+                do_sample=True,
+                eos_token_id=tokenizer.eos_token_id
+            )
+
+            # Decode the response
+            response = tokenizer.decode(output[0], skip_special_tokens=True)
+
+            # Display the response
+            st.text_area("Bot:", value=response, height=200, max_chars=None, key=None)
+
+        except Exception as e:
+            st.error(f"An error occurred: {e}")