import gradio as gr
from vllm import LLM, SamplingParams
from transformers import BitsAndBytesConfig

# Load quantized model
quant_config = BitsAndBytesConfig(load_in_4bit=True)
model_path = "ai4bharat/IndicTrans3-beta"  # Use quantized version if available
llm = LLM(model=model_path, quantization="bitsandbytes", quantization_config=quant_config)

def translate(text, src_lang, tgt_lang):
    prompt = f"Translate from {src_lang} to {tgt_lang}: {text}"
    sampling_params = SamplingParams(temperature=0.1, top_p=0.95, max_tokens=512)
    outputs = llm.generate([prompt], sampling_params)
    return outputs[0].outputs[0].text.strip()

iface = gr.Interface(
    fn=translate,
    inputs=[gr.Textbox(label="Input Text"), gr.Textbox(label="Source Language (e.g., en)"), gr.Textbox(label="Target Language (e.g., hi)")],
    outputs=gr.Textbox(label="Translated Text"),
    title="IndicTrans3 Translator"
)

iface.launch()