Update app.py

app.py

@@ -1,50 +1,49 @@
 import streamlit as st
-import os
-from langchain_groq import ChatGroq
-from dotenv import load_dotenv
 from transformers import AutoTokenizer, AutoModelForCausalLM
-
-# Load environment variables from .env file
-load_dotenv()
-groq_api_key = os.environ.get("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")
-
-# Initialize the Hugging Face model
-huggingface_model_name = "s0uL141/fine_tuned_science_gemma2b-it"  # Your Hugging Face model name
-tokenizer = AutoTokenizer.from_pretrained(huggingface_model_name)
-huggingface_model = AutoModelForCausalLM.from_pretrained(huggingface_model_name)
-
-# Initialize the ChatGroq model with the Hugging Face model name for inference
-llm = ChatGroq(
-    temperature=0,
-    model_name=huggingface_model_name,  # Use the Hugging Face model name here
-    api_key=groq_api_key
-)
-
-# Function to generate text using Hugging Face model (optional if you're using ChatGroq directly)
-def generate_response_huggingface(prompt, max_length=500):
+import torch
+
+# Cache the model loading to avoid reloading it on every interaction
+@st.cache_resource
+def load_model():
+    model_name = "s0uL141/Cyber_gemma2_2B_it"  # Replace with your Hugging Face repo or local path
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16)
+    return tokenizer, model
+
+# Load the model and tokenizer
+tokenizer, model = load_model()
+
+# Function to generate text based on the user prompt
+def generate_response(prompt, max_length=50):
+    # Tokenize input prompt
     inputs = tokenizer(prompt, return_tensors="pt")
-    output = huggingface_model.generate(inputs.input_ids, max_length=max_length)
+    # Generate response using the model
+    output = model.generate(inputs.input_ids, max_length=max_length, num_return_sequences=1)
+    # Decode the response and return
     return tokenizer.decode(output[0], skip_special_tokens=True)
 
 # Streamlit App
 def main():
-    st.title("Cybersecurity Q&A with ChatGroq and Hugging Face Model")
-    st.write("This app generates responses to your cybersecurity questions using ChatGroq with a Hugging Face model.")
+    # Set up the title and description for the app
+    st.title("Fine-Tuned Cyber Gemma 2b-it Model")
+    st.write("This app generates responses based on your input using a fine-tuned version of the Gemma 2b-it model.")
 
-    user_input = st.text_area("Enter your cybersecurity-related question below:", height=200)
+    # Text input area for the user to provide a prompt
+    user_input = st.text_area("Enter your prompt here:", height=200)
 
+    # Button to trigger text generation
     if st.button("Generate Response"):
+        # Check if user input is provided
         if user_input.strip() == "":
-            st.write("Please enter a valid question.")
+            st.write("Please enter a valid prompt.")
         else:
             with st.spinner("Generating response..."):
-                # Invoke the ChatGroq model with the user's input
-                response = llm.invoke(user_input)
-                
+                # Generate response using the model
+                response = generate_response(user_input)
                 # Display the generated response
                 st.write("### Model Response:")
                 st.write(response)
 
 # Entry point to run the app
 if __name__ == "__main__":
-    main()
+    main()