Update app.py

app.py

@@ -1,52 +1,32 @@
 import streamlit as st
+import io
 from PIL import Image
-from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM
 from easyocr import Reader
 
-# Initialize the image-to-text model
-image_to_text = pipeline("image-to-text", model="Salesforce/blip-image-captioning-large")
 
-# Load the OCR model and text explanation model
 ocr_reader = Reader(['en'])
-
-# Load the text explanation model
-text_generator = AutoModelForCausalLM.from_pretrained("HuggingFaceH4/zephyr-7b-alpha")
-text_tokenizer = AutoTokenizer.from_pretrained("HuggingFaceH4/zephyr-7b-alpha")
-
-# Define a function to extract text from an image using OCR
+text_generator = AutoModelForCausalLM.from_pretrained("gpt2")
+text_tokenizer = AutoTokenizer.from_pretrained("gpt2")
 def extract_text(image):
     return ocr_reader.readtext(image)
-
-# Define a function to explain the extracted text using text generation
-def explain_text(text, text_generator, text_tokenizer):
-    # Extracted text
-    extracted_text = " ".join([res[1] for res in text])
-
-    # Generate an explanation using the text explanation model
-    input_ids = text_tokenizer.encode(extracted_text, return_tensors="pt")
+def explain_text(text):
+    input_ids = text_tokenizer.encode(text, return_tensors="pt")
     explanation_ids = text_generator.generate(input_ids, max_length=100, num_return_sequences=1)
     explanation = text_tokenizer.decode(explanation_ids[0], skip_special_tokens=True)
-    
     return explanation
 
-# Create a Streamlit layout
-st.title("Text Extraction and Explanation")
-
-# Allow users to upload an image
+st.title("Text Classification Model")
 uploaded_file = st.file_uploader("Upload an image:")
 
-# Extract text from the uploaded image and explain it
 if uploaded_file is not None:
     image = Image.open(uploaded_file)
     ocr_results = extract_text(image)
-    image_caption = image_to_text(image)  # Use the image-to-text model
-    explanation = explain_text(ocr_results, text_generator, text_tokenizer)
-
+    extracted_text = " ".join([res[1] for res in ocr_results])
+    explanation = explain_text(extracted_text)
     st.markdown("**Extracted text:**")
-    st.markdown(" ".join([res[1] for res in ocr_results]))
-
-    st.markdown("**Image Caption:**")
-    st.markdown(image_caption[0]['caption'])  # Display the image caption
+    st.markdown(extracted_text)
 
     st.markdown("**Explanation:**")
     st.markdown(explanation)