Spaces:
Running
Running
| import gradio as gr | |
| import torch | |
| from transformers import AutoTokenizer, AutoModelForCausalLM | |
| from huggingface_hub import login | |
| import os | |
| # ----------------------------- | |
| # 1. Login using HF token (from environment variable) | |
| # ----------------------------- | |
| hf_token = os.environ.get("HF_TOKEN") | |
| if hf_token: | |
| login(token=hf_token) | |
| else: | |
| print("⚠️ Warning: HF_TOKEN not set. Add it in the Space settings.") | |
| # ----------------------------- | |
| # 2. Load Model | |
| # ----------------------------- | |
| model_id = "ibm-granite/granite-3.3-2b-instruct" | |
| tokenizer = AutoTokenizer.from_pretrained(model_id, use_auth_token=True) | |
| model = AutoModelForCausalLM.from_pretrained( | |
| model_id, | |
| torch_dtype=torch.float16, | |
| device_map="auto", | |
| use_auth_token=True | |
| ) | |
| # ----------------------------- | |
| # 3. Functions | |
| # ----------------------------- | |
| def classify_sentiment(text): | |
| prompt = f""" | |
| Classify the sentiment of the following review as one word: Positive, Negative, or Neutral. | |
| Review: "{text}" | |
| Only respond with the sentiment label. | |
| """ | |
| inputs = tokenizer(prompt, return_tensors="pt").to(model.device) | |
| output = model.generate( | |
| **inputs, | |
| max_new_tokens=5, | |
| temperature=0.1, | |
| repetition_penalty=1.5 | |
| ) | |
| result = tokenizer.decode(output[0], skip_special_tokens=True) | |
| final = result.split()[-1].strip() | |
| return final | |
| def summarize_text(text): | |
| prompt = f""" | |
| Summarize the following text in 3–5 bullet points. | |
| Make the summary short, simple, and avoid repeating sentences. | |
| Text: {text} | |
| Summary: | |
| """ | |
| inputs = tokenizer(prompt, return_tensors="pt").to(model.device) | |
| output = model.generate( | |
| **inputs, | |
| max_new_tokens=150, | |
| temperature=0.7, | |
| top_p=0.9, | |
| repetition_penalty=2.0, | |
| ) | |
| result = tokenizer.decode(output[0], skip_special_tokens=True) | |
| summary = result.split("Summary:")[-1].strip() | |
| return summary | |
| # ----------------------------- | |
| # 4. Gradio Interface | |
| # ----------------------------- | |
| with gr.Blocks() as demo: | |
| gr.Markdown("## 🟦 Granite 3.3-2B — Sentiment + Summarization Demo") | |
| with gr.Tab("Sentiment Classification"): | |
| inp = gr.Textbox(label="Enter Review", lines=4) | |
| out = gr.Textbox(label="Predicted Sentiment", lines=1) | |
| btn = gr.Button("Classify") | |
| btn.click(classify_sentiment, inp, out) | |
| with gr.Tab("Summarization"): | |
| inp2 = gr.Textbox(label="Enter Text to Summarize", lines=8) | |
| out2 = gr.Textbox(label="Summary", lines=8) | |
| btn2 = gr.Button("Summarize") | |
| btn2.click(summarize_text, inp2, out2) | |
| demo.launch() |