Update app.py

app.py

@@ -2,9 +2,9 @@ import torch
 import streamlit as st
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
-# Load the model and tokenizer
-model_name = "nvidia/Llama-3.1-Nemotron-70B-Instruct-HF"
-model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.bfloat16, device_map="auto")
+# Use a smaller model for testing; switch to the larger model when GPU resources are available
+model_name = "gpt2"
+model = AutoModelForCausalLM.from_pretrained(model_name)
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 
 # Set up the Streamlit app interface
@@ -20,8 +20,8 @@ if st.button("Generate Response"):
         with st.spinner("Generating response..."):
             # Tokenize input
             tokenized_prompt = tokenizer(prompt, return_tensors="pt")
-            input_ids = tokenized_prompt['input_ids'].cuda() if torch.cuda.is_available() else tokenized_prompt['input_ids']
-            attention_mask = tokenized_prompt['attention_mask'].cuda() if torch.cuda.is_available() else tokenized_prompt['attention_mask']
+            input_ids = tokenized_prompt['input_ids']
+            attention_mask = tokenized_prompt['attention_mask']
 
             # Generate the model's response
             response_token_ids = model.generate(