|
|
import streamlit as st |
|
|
from transformers import MBartForConditionalGeneration, MBart50TokenizerFast |
|
|
|
|
|
@st.cache_resource |
|
|
def load_model(): |
|
|
tokenizer = MBart50TokenizerFast.from_pretrained("MahmutCanBoran/mbart-audi-diagnosis-agent") |
|
|
model = MBartForConditionalGeneration.from_pretrained("MahmutCanBoran/mbart-audi-diagnosis-agent") |
|
|
return tokenizer, model |
|
|
|
|
|
tokenizer, model = load_model() |
|
|
|
|
|
st.title("🔧 Audi AI Diagnosis Agent") |
|
|
st.markdown("Enter your Audi issue in **English**, and the AI will try to diagnose it.") |
|
|
|
|
|
text = st.text_area("🔍 What's the problem?", "") |
|
|
|
|
|
if st.button("🧠 Diagnose"): |
|
|
if not text.strip(): |
|
|
st.warning("Please enter a vehicle problem.") |
|
|
else: |
|
|
inputs = tokenizer( |
|
|
text, |
|
|
return_tensors="pt", |
|
|
truncation=True, |
|
|
padding="max_length", |
|
|
max_length=128 |
|
|
).to(model.device) |
|
|
|
|
|
output_ids = model.generate( |
|
|
**inputs, |
|
|
max_length=128, |
|
|
num_beams=4, |
|
|
early_stopping=True, |
|
|
forced_bos_token_id=tokenizer.lang_code_to_id["en_XX"] |
|
|
) |
|
|
|
|
|
result = tokenizer.decode(output_ids[0], skip_special_tokens=True) |
|
|
st.success("🔧 AI Diagnosis:") |
|
|
st.markdown(f"**{result}**") |
