Update app.py

app.py

@@ -19,8 +19,8 @@ ocr_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
 ocr_model = AutoModelForCausalLM.from_pretrained("microsoft/Florence-2-large", torch_dtype=ocr_dtype, trust_remote_code=True).to(ocr_device)
 ocr_processor = AutoProcessor.from_pretrained("microsoft/Florence-2-large", trust_remote_code=True)
 
-# Load GPT-2 Large model for text generation
-gpt2_generator = pipeline('text-generation', model='gpt2-large')
+# Load GPT-2 (124M) model for text generation
+gpt2_generator = pipeline('text-generation', model='gpt2')
 
 # Load speaker embedding
 embeddings_dataset = load_dataset("Matthijs/cmu-arctic-xvectors", split="validation")
@@ -43,7 +43,7 @@ def process_image(image):
         )
         extracted_text = ocr_processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
 
-        # Generate context using GPT-2 Large
+        # Generate context using GPT-2 (124M)
         prompt = f"Determine the context of this image based on the caption and extracted text. Caption: {caption}. Extracted text: {extracted_text}. Context:"
         context_output = gpt2_generator(prompt, max_length=100, num_return_sequences=1)
         context = context_output[0]['generated_text']
@@ -75,8 +75,9 @@ iface = gr.Interface(
         gr.Textbox(label="Extracted Text (OCR)"),
         gr.Textbox(label="Generated Context")
     ],
-    title="SeeSay Contextualizer with GPT-2 Large",
-    description="Upload an image to generate a caption, extract text, create audio from context, and determine the context using GPT-2 Large."
+    title="SeeSay Contextualizer with GPT-2 (124M)",
+    description="Upload an image to generate a caption, extract text, create audio from context, and determine the context using GPT-2."
 )
 
 iface.launch()
+