Spaces:
Sleeping
Sleeping
File size: 1,390 Bytes
19ae9ab 2645340 1fa58c0 2645340 f6a4ebc 2645340 f6a4ebc 2645340 03444e8 19ae9ab f6a4ebc 03444e8 f6a4ebc 4005d9f f6a4ebc dc59162 f6a4ebc 4005d9f 0e8a501 4005d9f dc59162 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 |
import streamlit as st
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
# Load the model and tokenizer with memory optimizations
model_name = "Tom158/Nutri_Assist"
model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16, device_map="auto")
tokenizer = AutoTokenizer.from_pretrained(model_name)
# Set pad token if not already set
if model.config.pad_token_id is None:
model.config.pad_token_id = model.config.eos_token_id
# Streamlit App Interface
st.title("Nutrition Chatbot")
user_input = st.text_input("Ask me about nutrition:")
if user_input:
# Truncate input and convert to tensors
inputs = tokenizer.encode_plus(user_input, return_tensors="pt", padding=True, truncation=True, max_length=512)
input_ids = inputs['input_ids']
attention_mask = inputs['attention_mask']
# Generate output with attention mask and pad token ID
try:
# Limit output length to save memory
outputs = model.generate(input_ids, attention_mask=attention_mask, max_length=100,
temperature=0.7, top_k=50, num_return_sequences=1)
# Decode the output and display
decoded_output = tokenizer.decode(outputs[0], skip_special_tokens=True)
st.write("Decoded Answer:", decoded_output)
except Exception as e:
st.write("Error generating output:", str(e))
|