import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

MODEL = "Suvir-Misra/GST-Cases-Laws"
tokenizer = AutoTokenizer.from_pretrained(MODEL)

# CPU-ONLY: No quantization, pure float16
model = AutoModelForCausalLM.from_pretrained(
    MODEL,
    torch_dtype=torch.float16,
    low_cpu_mem_usage=True,
    device_map="cpu"
)

def generate_gst_text(prompt):
    inputs = tokenizer(prompt[:300], return_tensors="pt")
    output = model.generate(
        **inputs,
        max_new_tokens=80,        # Keep low for speed
        do_sample=False,          # Greedy = 3x faster
        pad_token_id=tokenizer.eos_token_id
    )
    return tokenizer.decode(output[0][inputs.input_ids.shape[1]:], skip_special_tokens=True)

gr.Interface(fn=generate_gst_text, inputs=gr.Textbox(placeholder="GST query..."), outputs="text").launch()