Update app.py

app.py

@@ -1,5 +1,5 @@
 import nltk
-from transformers import VisionEncoderDecoderModel, AutoTokenizer, ViTImageProcessor, GPTNeoForCausalLM, GPT2Tokenizer
+from transformers import VisionEncoderDecoderModel, AutoTokenizer, ViTImageProcessor, pipeline
 import torch
 from PIL import Image
 import streamlit as st
@@ -7,6 +7,7 @@ from nltk.corpus import stopwords
 import os
 from io import BytesIO
 
+
 os.system('pip install --upgrade transformers')
 # os.system('pip install nltk')
 nltk.download('stopwords')
@@ -24,9 +25,8 @@ feature_extractor = ViTImageProcessor.from_pretrained(
 tokenizer = AutoTokenizer.from_pretrained(
     'nlpconnect/vit-gpt2-image-captioning')
 
-# Load GPT-2 model and tokenizer
-gpt2_model = GPTNeoForCausalLM.from_pretrained('EleutherAI/gpt-neo-1.3-125')
-gpt2_tokenizer = GPT2Tokenizer.from_pretrained('EleutherAI/gpt-neo-1.3-125')
+# Set up text generation pipeline
+generator = pipeline('text-generation', model='EleutherAI/gpt-neo-2.7B')
 
 # Device configuration
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
@@ -115,16 +115,11 @@ def main():
         # Prepare the prompt
         prompt = f"You are a knowledgeable assistant that provides nutritional advice based on a list of ingredients. The identified ingredients are: {preds_str}. Note that some ingredients may not make sense, so use the ones that do. Can you provide a nutritional analysis and suggestions for improvement?"
 
-        # Encode and add special tokens
-        input_ids = gpt2_tokenizer.encode(prompt, return_tensors='pt')
-
-        # Generate a sequence of text
-        output = gpt2_model.generate(
-            input_ids, max_length=200, temperature=0.7, pad_token_id=gpt2_tokenizer.eos_token_id)
+         # Generate a sequence of text
+        suggestions = generator(prompt, do_sample=True, min_length=200)
 
-        # Decode the output
-        suggestions = gpt2_tokenizer.decode(
-            output[:, input_ids.shape[-1]:][0], clean_up_tokenization_spaces=True)
+        # Extract the generated text
+        suggestions = suggestions[0]['generated_text'][len(prompt):]
 
         st.subheader("Nutritional Analysis and Suggestions:")
         st.write(suggestions)