Spaces:
Sleeping
Sleeping
File size: 3,254 Bytes
f3a781a |
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 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 |
import streamlit as st
from transformers import pipeline, AutoTokenizer, AutoModelForQuestionAnswering
from PIL import Image
import easyocr
import os
from groq import Groq, APIConnectionError, AuthenticationError
# OCR Function
def extract_text_from_image(image):
reader = easyocr.Reader(['en'])
result = reader.readtext(image)
extracted_text = " ".join([detection[1] for detection in result])
return extracted_text
# Question Answering Function (DistilBERT)
@st.cache_resource
def load_qa_model():
model_name = "distilbert/distilbert-base-cased-distilled-squad"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForQuestionAnswering.from_pretrained(model_name)
nlp = pipeline('question-answering', model=model, tokenizer=tokenizer)
return nlp
def answer_question(context, question, qa_model):
result = qa_model({'question': question, 'context': context})
return result['answer']
# Groq API Function
def groq_chat(prompt):
try:
client = Groq(api_key=os.environ.get("GROQ_API_KEY"))
chat_completion = client.chat.completions.create(
messages=[{"role": "user", "content": prompt}],
model="llama-3.3-70b-versatile",
)
return chat_completion.choices[0].message.content
except APIConnectionError as e:
return f"Groq API Connection Error: {e}"
except AuthenticationError as e:
return f"Groq API Authentication Error: {e}"
except Exception as e:
return f"General Groq API Error: {e}"
# Streamlit App
def main():
st.title("Image Text & Question Answering Chatbot")
uploaded_file = st.file_uploader("Upload an image", type=["jpg", "jpeg", "png"])
if uploaded_file is not None:
image = Image.open(uploaded_file)
st.image(image, caption="Uploaded Image", use_container_width=True)
if st.button("Extract Text and Enable Question Answering"):
with st.spinner("Extracting text..."):
extracted_text = extract_text_from_image(image)
st.session_state.extracted_text = extracted_text # Store in session state
st.write("Extracted Text:")
st.write(st.session_state.extracted_text)
if "extracted_text" in st.session_state: # Check if extracted_text is in session state
qa_model = load_qa_model()
question = st.text_input("Ask a question about the image text:")
if st.button("Answer"):
if question:
with st.spinner("Answering..."):
answer = answer_question(st.session_state.extracted_text, question, qa_model)
st.write("Answer:", answer)
else:
st.warning("Please enter a question.")
# Groq Chat Section
st.subheader("General Chat (Powered by Groq)")
groq_prompt = st.text_input("Enter your message:")
if st.button("Send"):
if groq_prompt:
with st.spinner("Generating response..."):
groq_response = groq_chat(groq_prompt)
st.write("Response:", groq_response)
else:
st.warning("Please enter a message.")
if __name__ == "__main__":
main() |