import gradio as gr
from transformers import AutoModelForSequenceClassification, AutoTokenizer

# Load your fine-tuned model and tokenizer
model_name = "quadranttechnologies/Receipt_Image_Analyzer"
model = AutoModelForSequenceClassification.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)

# Define a prediction function
def analyze_receipt(receipt_text):
    inputs = tokenizer(receipt_text, return_tensors="pt", truncation=True, padding=True)
    outputs = model(**inputs)
    logits = outputs.logits
    predicted_class = logits.argmax(-1).item()
    return f"Predicted Class: {predicted_class}"

# Create a Gradio interface
interface = gr.Interface(
    fn=analyze_receipt,
    inputs="text",
    outputs="text",
    title="Receipt Image Analyzer",
    description="Analyze receipts for relevant information using a fine-tuned LLM model.",
)

# Launch the Gradio app
interface.launch()