voice-to-code / app.py
donnamae's picture
Update app.py
e0c7921 verified
import torch
import librosa
import gradio as gr
from transformers import WhisperProcessor, WhisperForConditionalGeneration, AutoTokenizer, AutoModelForCausalLM
# Load models from the Space or from Hugging Face Hub
whisper_model = WhisperForConditionalGeneration.from_pretrained("donnamae/whisper-finetuned-cebuano-accent", token=True)
whisper_processor = WhisperProcessor.from_pretrained("donnamae/whisper-finetuned-cebuano-accent", token=True)
code_tokenizer = AutoTokenizer.from_pretrained("meta-llama/CodeLlama-7b-Instruct-hf")
code_model = AutoModelForCausalLM.from_pretrained(
"meta-llama/CodeLlama-7b-Instruct-hf",
torch_dtype="auto",
device_map="auto",
trust_remote_code=True
).to("cuda" if torch.cuda.is_available() else "cpu")
def transcribe_and_generate(audio):
audio_data, sr = librosa.load(audio, sr=16000)
input_features = whisper_processor(audio_data, sampling_rate=sr, return_tensors="pt").input_features
predicted_ids = whisper_model.generate(input_features)
transcription = whisper_processor.batch_decode(predicted_ids, skip_special_tokens=True)[0]
# Format prompt for code generation
prompt = f"# Task: {transcription.strip()}\n\n```python\n"
inputs = code_tokenizer(prompt, return_tensors="pt").to(code_model.device)
outputs = code_model.generate(**inputs, max_length=256)
generated_text = code_tokenizer.decode(outputs[0], skip_special_tokens=True)
# Extract clean code
generated_code = generated_text.replace(prompt, "").strip().split("```")[0]
return transcription, generated_code
demo = gr.Interface(
fn=transcribe_and_generate,
inputs=gr.Audio(type="filepath"),
outputs=[gr.Text(label="Transcribed Command"), gr.Code(label="Generated Code")],
title="Voice-to-Code Generator",
description="Speak your coding command. The system will transcribe and generate the corresponding code."
)
demo.launch(share=True)