import os os.environ["STREAMLIT_HOME"] = os.getcwd() os.environ["STREAMLIT_RUNTIME_METRICS_ENABLED"] = "false" os.makedirs(".streamlit", exist_ok=True) import streamlit as st import torch import joblib import numpy as np import random from PIL import Image from transformers import AutoTokenizer, AutoModel, ViTModel, ViTImageProcessor # CPU device only device = torch.device("cpu") # Define Swahili VQA Model class SwahiliVQAModel(torch.nn.Module): def __init__(self, num_answers): super().__init__() self.vision_encoder = ViTModel.from_pretrained('google/vit-base-patch16-224-in21k') self.text_encoder = AutoModel.from_pretrained("benjamin/roberta-base-wechsel-swahili") self.fusion = torch.nn.Sequential( torch.nn.Linear(768 + 768, 512), torch.nn.ReLU(), torch.nn.Dropout(0.3), torch.nn.LayerNorm(512) ) self.classifier = torch.nn.Linear(512, num_answers) def forward(self, image, input_ids, attention_mask): vision_outputs = self.vision_encoder(pixel_values=image) image_feats = vision_outputs.last_hidden_state[:, 0, :] text_outputs = self.text_encoder(input_ids=input_ids, attention_mask=attention_mask) text_feats = text_outputs.last_hidden_state[:, 0, :] combined = torch.cat([image_feats, text_feats], dim=1) fused = self.fusion(combined) return self.classifier(fused) # Load label encoder le = joblib.load("Vit_3895_label_encoder_best.pkl") # Load model weights normally โ no override model = SwahiliVQAModel(num_answers=len(le.classes_)).to(device) state_dict = torch.load("Vit_3895_best_model_epoch25.pth", map_location=device) model.load_state_dict(state_dict) model.eval() # Load tokenizer and processor tokenizer = AutoTokenizer.from_pretrained("benjamin/roberta-base-wechsel-swahili") vit_processor = ViTImageProcessor.from_pretrained('google/vit-base-patch16-224-in21k') # Streamlit UI st.set_page_config(page_title="Swahili VQA", layout="wide") st.title("๐ฆ Swahili Visual Question Answering (VQA)") uploaded_image = st.file_uploader("๐ Pakia picha hapa:", type=["jpg", "jpeg", "png"]) def generate_random_color(): return f"rgb({random.randint(150, 255)}, {random.randint(80, 200)}, {random.randint(80, 200)})" col1, col2 = st.columns([1, 2], gap="large") with col1: if uploaded_image: st.image(uploaded_image, caption="Picha Iliyopakiwa", use_container_width=True) st.markdown("
", unsafe_allow_html=True) with col2: st.markdown("