Update app.py

app.py

@@ -1,49 +1,57 @@
 import streamlit as st
-from transformers import RagTokenizer, RagRetriever, RagSequenceForGeneration
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 import torch
 
-# Load the fine-tuned FLAN-T5 model and tokenizer
-rag_tokenizer = RagTokenizer.from_pretrained("fine_tuned_flan_t5")  # Your fine-tuned model path
-rag_retriever = RagRetriever.from_pretrained("facebook/rag-token-nq")  # Pre-trained retriever
-rag_model = RagSequenceForGeneration.from_pretrained("Izza-shahzad-13/fine-tuned-flan-t5")  # Your fine-tuned model
+# Hugging Face Token for Authentication
+HUGGINGFACE_TOKEN = "your_hugging_face_token_here"  # Replace with your token
 
-# Setup device
+# Function to load model and tokenizer (local or Hugging Face with token)
+def load_model(model_path):
+    try:
+        tokenizer = AutoTokenizer.from_pretrained(model_path, use_auth_token=HUGGINGFACE_TOKEN)
+        model = AutoModelForSeq2SeqLM.from_pretrained(model_path, use_auth_token=HUGGINGFACE_TOKEN)
+        return tokenizer, model
+    except Exception as e:
+        st.error(f"Error loading model: {e}")
+        return None, None
+
+# Set device (use GPU if available)
 device = "cuda" if torch.cuda.is_available() else "cpu"
-rag_model.to(device)
-
-# Function to generate RAG response
-def generate_rag_response(input_text):
-    # Tokenize the input
-    inputs = rag_tokenizer(input_text, return_tensors="pt", padding=True, truncation=True)
-
-    # Retrieve relevant documents based on input
-    input_ids = inputs['input_ids'].to(device)
-    retrieved_docs = rag_retriever(input_ids)
-
-    # Generate the response from the retrieved context
-    generated_ids = rag_model.generate(
-        input_ids=input_ids,
-        context_input_ids=retrieved_docs['context_input_ids'].to(device),
-        context_attention_mask=retrieved_docs['context_attention_mask'].to(device),
-        max_length=200,
-        num_beams=4,
-        top_p=0.9,
-        top_k=50,
-        temperature=0.7,
-        no_repeat_ngram_size=3,
-        early_stopping=True
-    )
-    
-    # Decode the generated response
-    return rag_tokenizer.decode(generated_ids[0], skip_special_tokens=True)
-
-# Streamlit UI setup
-st.title("Mental Health Counseling Assistant with RAG")
-
-# Input for user query
+
+# Path to your model (either a local path or a Hugging Face model name)
+model_path = "Izza-shahzad-13/fine-tuned-flan-t5"  # Use your Hugging Face model identifier
+
+# Load tokenizer and model
+tokenizer, model = load_model(model_path)
+if model:
+    model.to(device)
+
+# Function to generate response from the model
+def generate_response(input_text):
+    inputs = tokenizer(input_text, return_tensors="pt", padding=True, truncation=True, max_length=512).to(device)
+    with torch.no_grad():
+        output = model.generate(
+            inputs['input_ids'],
+            max_length=500,
+            num_beams=4,
+            top_p=0.9,
+            top_k=50,
+            temperature=0.7,
+            do_sample=True,
+            no_repeat_ngram_size=3,
+            early_stopping=True
+        )
+    response = tokenizer.decode(output[0], skip_special_tokens=True)
+    return response
+
+# Streamlit app interface
+st.title("FLAN-T5 Mental Health Counseling Assistant")
+st.write("Type your thoughts or feelings, and let the model respond.")
+
+# User input for interaction
 user_input = st.text_input("How are you feeling today?")
 
+# Generate and display model response when input is provided
 if user_input:
-    # Generate and display the response using the RAG model
-    response = generate_rag_response(user_input)
+    response = generate_response(user_input)
     st.write("Model Response:", response)