app.py

import gradio as gr
from transformers import pipeline
from langchain import PromptTemplate
from langchain.document_loaders import WebBaseLoader
from langchain.schema import StrOutputParser
from langchain.schema.prompt_template import format_document
from langchain.chains import LLMChain
from langchain_google_genai import ChatGoogleGenerativeAI
import os
from PIL import Image
import json

# Retrieve the API keys and other secrets from the environment
api_key = os.environ.get('GOOGLE_API_KEY')
if api_key is None:
    raise ValueError("No API key found. Please set the 'GOOGLE_API_KEY' environment variable.")

tracking_id = os.environ.get('TRACKING_ID')
if tracking_id is None:
    raise ValueError("No tracking ID found. Please set the 'TRACKING_ID' environment variable.")

initial_prompt = os.environ.get('initial_prompt')
if initial_prompt is None:
    raise ValueError("No initial prompt found. Please set the 'initial_prompt' environment variable.")

description_json = os.environ.get('description')
if description_json is None:
    raise ValueError("No description found. Please set the 'description' environment variable.")

# Convert the description JSON string to a dictionary
description = json.loads(description_json)

# Set the API key for Google
os.environ['GOOGLE_API_KEY'] = api_key

# Initialize the OCR pipeline
ocr_pipe = pipeline("image-to-text", model="jinhybr/OCR-Donut-CORD")

# Initialize the LLM
llm_model = ChatGoogleGenerativeAI(model="gemini-pro", temperature=0.4, top_p=0.85)

# Define the prompt template
prompt = PromptTemplate(input_variables=['task_type', 'task_number', 'question', 'content', 'description'], template=initial_prompt)

# Define the LLM chain
chain = LLMChain(llm=llm_model, prompt=prompt)

def evaluate(task_type, task_number, question, image):
    # Ensure the image is in the correct format
    if isinstance(image, str):
        # Load the image if it's a URL or path
        image = Image.open(image)
    
    # Process the image to extract text
    text_content = ocr_pipe(image)
    content = text_content[0]['generated_text']

    # Run the chain
    result = chain.run({
        'task_type': task_type,
        'task_number': task_number,
        'question': question,
        'content': content,
        'description': description[(task_type, task_number)]
    })

    return result

# Create the Gradio interface
inputs = [
    gr.Dropdown(choices=["Academic", "General"], label="Test Type", value="Academic"),
    gr.Dropdown(choices=["Task 1", "Task 2"], label="Task Number", value="Task 1"),
    gr.Textbox(label="Question", value=""),
    gr.Image(type="pil", label="Upload Image")
]

footer_html_with_analytics = f"""
<script async src="https://www.googletagmanager.com/gtag/js?id={tracking_id}"></script>
<script>
  window.dataLayer = window.dataLayer || [];
  function gtag(){{dataLayer.push(arguments);}}
  gtag('js', new Date());
  gtag('config', '{tracking_id}');
</script>
<link rel="stylesheet" href="https://cdnjs.cloudflare.com/ajax/libs/font-awesome/6.0.0-beta3/css/all.min.css">
<div style='text-align: center; margin-top: 20px;'>
    <p>Developed by Hossein Mohseni</p>
    <p>Contact Information:</p>
    <p>
        <a href='mailto:[email protected]' style='margin-right: 10px;'>
            <i class='fas fa-envelope'></i>
        </a>
        <a href='https://www.linkedin.com/in/mohsenihossein/' target='_blank' style='margin-right: 10px;'>
            <i class='fab fa-linkedin'></i>
        </a>
        <a href='https://t.me/mohsenih1999' target='_blank'>
            <i class='fab fa-telegram'></i>
        </a>
    </p>
    <p>This application is a demonstration. To enhance and improve it, your feedback is highly appreciated.</p>
</div>
"""

outputs = gr.Markdown(label="Result")

# Define the Gradio Blocks and Interface
with gr.Blocks() as demo:
    gr.Markdown("# IELTS Writing Evaluation")
    gr.Interface(fn=evaluate, inputs=inputs, outputs=outputs)
    gr.HTML(footer_html_with_analytics)

# Launch the interface
demo.launch(share=True, debug=True)