app.py

import gradio as gr
from transformers import AutoTokenizer, AutoConfig, AutoModelForSequenceClassification

# Load model and tokenizer
model_name = "cheberle/autotrain-35swc-b4r9z"
tokenizer = AutoTokenizer.from_pretrained(model_name)

# Explicitly define the model configuration if needed
config = AutoConfig.from_pretrained(model_name)
model = AutoModelForSequenceClassification.from_pretrained(model_name, config=config)

# Inference function
def classify_text(input_text):
    inputs = tokenizer(input_text, return_tensors="pt", padding=True, truncation=True)
    outputs = model(**inputs)
    probabilities = outputs.logits.softmax(dim=-1).tolist()[0]
    labels = {i: f"Label {i}" for i in range(len(probabilities))}  # Define label mapping if needed
    result = {labels[i]: prob for i, prob in enumerate(probabilities)}
    return result

# Gradio interface
interface = gr.Interface(
    fn=classify_text,
    inputs="text",
    outputs="label",
    title="DeepSeek-R1 Text Classification",
    description="Classify text inputs using the DeepSeek-R1 model."
)

# Launch the app
if __name__ == "__main__":
    interface.launch()