Spaces:
Sleeping
Sleeping
| import streamlit as st | |
| from transformers import AutoTokenizer, AutoModelForSeq2SeqLM | |
| import torch | |
| import os | |
| # Retrieve the Hugging Face token from environment variable | |
| HUGGINGFACE_TOKEN = os.getenv("HF_ACCESS_TOKEN") | |
| # Check if the token is available | |
| if not HUGGINGFACE_TOKEN: | |
| st.error("Hugging Face token not found. Please set the HUGGINGFACE_TOKEN environment variable.") | |
| st.stop() # Stop the app if the token is not found | |
| # Function to load model and tokenizer (local or Hugging Face with token) | |
| def load_model(model_path): | |
| try: | |
| # Load tokenizer and model with authentication token if required | |
| 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 = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
| # 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 and tokenizer: | |
| model.to(device) | |
| else: | |
| st.stop() # Stop the app if model or tokenizer failed to load | |
| # 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_area("How are you feeling today?", placeholder="Type here...") | |
| # Generate and display model response when input is provided | |
| if user_input.strip(): # Check if input is not empty | |
| with st.spinner("Generating response..."): | |
| response = generate_response(user_input) | |
| st.write("Model Response:", response) | |
| else: | |
| st.info("Please enter your thoughts or feelings in the text area above.") |