import gradio as gr
import torchaudio
import torch
from transformers import Wav2Vec2Processor, Wav2Vec2ForCTC

# Load model directly from Hugging Face Hub
model_id = "onepunchgin/ASR-Model"

processor = Wav2Vec2Processor.from_pretrained(model_id)
model = Wav2Vec2ForCTC.from_pretrained(model_id)
model.eval()

def transcribe(audio):
    waveform, sr = torchaudio.load(audio)
    if sr != 16000:
        waveform = torchaudio.transforms.Resample(orig_freq=sr, new_freq=16000)(waveform)
    inputs = processor(waveform.squeeze().numpy(), sampling_rate=16000, return_tensors="pt", padding=True)
    with torch.no_grad():
        logits = model(**inputs).logits
    predicted_ids = torch.argmax(logits, dim=-1)
    return processor.batch_decode(predicted_ids)[0]

iface = gr.Interface(fn=transcribe,
                     inputs=gr.Audio(type="filepath"),
                     outputs="text",
                     title="Kannada Speech Recognition",
                     description="Upload a Kannada audio file to transcribe using CCC-wav2vec model")

iface.launch()