f

app.py

@@ -1,84 +1,55 @@
 import gradio as gr
 import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import AutoTokenizer, AutoModelForSequenceClassification, pipeline
+
+MODEL_NAME = "cheberle/autotrain-35swc-b4r9z"
 
 # ---------------------------------------------------------------------------
-# 1) Load the model and tokenizer
+# 1) Load the tokenizer and model for sequence classification
 # ---------------------------------------------------------------------------
-# If you want to load in 8-bit or 4-bit precision with bitsandbytes, 
-# uncomment and install bitsandbytes, and set load_in_8bit=True or load_in_4bit=True.
-# For example:
-#
-# from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
-# bnb_config = BitsAndBytesConfig(
-#     load_in_4bit=True,    # or load_in_8bit=True
-#     bnb_4bit_compute_dtype=torch.float16,  # recommended compute dtype
-#     bnb_4bit_use_double_quant=True,        # optional
-#     bnb_4bit_quant_type='nf4',             # optional
-# )
-#
-# model = AutoModelForCausalLM.from_pretrained(
-#     "cheberle/autotrain-35swc-b4r9z",
-#     quantization_config=bnb_config,
-#     device_map="auto",
-#     trust_remote_code=True
-# )
-# tokenizer = AutoTokenizer.from_pretrained("cheberle/autotrain-35swc-b4r9z", trust_remote_code=True)
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True)
+model = AutoModelForSequenceClassification.from_pretrained(MODEL_NAME, trust_remote_code=True)
 
-# For a standard FP16 or FP32 load (no bitsandbytes):
-model = AutoModelForCausalLM.from_pretrained(
-    "cheberle/autotrain-35swc-b4r9z",
-    torch_dtype=torch.float16,    # Or "auto", or float32
-    trust_remote_code=True
-)
-tokenizer = AutoTokenizer.from_pretrained(
-    "cheberle/autotrain-35swc-b4r9z",
-    trust_remote_code=True
-)
+# Create a pipeline for text classification
+classifier = pipeline("text-classification", model=model, tokenizer=tokenizer)
 
 # ---------------------------------------------------------------------------
-# 2) Define a text generation function
+# 2) Define inference function
 # ---------------------------------------------------------------------------
-def generate_text(prompt):
-    # Tokenize input
-    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
-
-    # Generate output (configure generation args as needed)
-    with torch.no_grad():
-        outputs = model.generate(
-            **inputs,
-            max_new_tokens=128,
-            temperature=0.7,
-            top_p=0.9,
-            do_sample=True
-        )
-
-    # Decode
-    decoded = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return decoded
+def classify_text(text):
+    """
+    Return the classification results in the format:
+    [
+      {
+        'label': 'POSITIVE',
+        'score': 0.98
+      }
+    ]
+    """
+    results = classifier(text)
+    return results
 
 # ---------------------------------------------------------------------------
-# 3) Create the Gradio interface
+# 3) Build the Gradio UI
 # ---------------------------------------------------------------------------
 with gr.Blocks() as demo:
-    gr.Markdown("<h3>Demo: cheberle/autotrain-35swc-b4r9z</h3>")
+    gr.Markdown("<h3>Text Classification Demo</h3>")
     
     with gr.Row():
-        with gr.Column():
-            prompt_in = gr.Textbox(
-                lines=5, 
-                label="Enter your prompt",
-                placeholder="Ask something here..."
-            )
-            submit_btn = gr.Button("Generate")
-        with gr.Column():
-            output_box = gr.Textbox(lines=15, label="Model Output")
+        input_text = gr.Textbox(
+            lines=3, 
+            label="Enter text to classify",
+            placeholder="Type something..."
+        )
+        output = gr.JSON(label="Classification Output")
+
+    classify_btn = gr.Button("Classify")
 
-    # Define what happens on button click
-    submit_btn.click(fn=generate_text, inputs=prompt_in, outputs=output_box)
+    # Link the button to the function
+    classify_btn.click(fn=classify_text, inputs=input_text, outputs=output)
 
 # ---------------------------------------------------------------------------
-# 4) Launch!
+# 4) Launch the demo
 # ---------------------------------------------------------------------------
 if __name__ == "__main__":
     demo.launch()