Spaces:
Sleeping
Sleeping
| # -*- coding: utf-8 -*- | |
| """Untitled1.ipynb | |
| Automatically generated by Colab. | |
| Original file is located at | |
| https://colab.research.google.com/drive/1vA1O3q8yuzV8Hi3O8LhNuLGWS18yVEkb | |
| """ | |
| import streamlit as st | |
| import PIL.Image | |
| import base64 | |
| import time | |
| import os | |
| import torch | |
| from transformers import AutoProcessor, LlavaForConditionalGeneration, BitsAndBytesConfig | |
| from peft import PeftModel | |
| from deep_translator import GoogleTranslator | |
| def load_model(): | |
| model_id = "HuggingFaceH4/vsft-llava-1.5-7b-hf-trl" | |
| quantization_config = BitsAndBytesConfig(load_in_4bit=True) | |
| base_model = LlavaForConditionalGeneration.from_pretrained(model_id, quantization_config=quantization_config, torch_dtype=torch.float16) | |
| # Load the PEFT Lora adapter | |
| peft_lora_adapter_path = "Praveen0309/llava-1.5-7b-hf-ft-mix-vsft-3" | |
| peft_lora_adapter = PeftModel.from_pretrained(base_model, peft_lora_adapter_path, adapter_name="lora_adapter") | |
| base_model.load_adapter(peft_lora_adapter_path, adapter_name="lora_adapter") | |
| processor = AutoProcessor.from_pretrained("HuggingFaceH4/vsft-llava-1.5-7b-hf-trl") | |
| return base_model, processor | |
| base_model, processor = load_model() | |
| # Function to translate text from Bengali to English | |
| def deep_translator_bn_en(input_sentence): | |
| english_translation = GoogleTranslator(source="bn", target="en").translate(input_sentence) | |
| return english_translation | |
| # Function to translate text from English to Bengali | |
| def deep_translator_en_bn(input_sentence): | |
| bengali_translation = GoogleTranslator(source="en", target="bn").translate(input_sentence) | |
| return bengali_translation | |
| def inference(image, image_prompt): | |
| prompt = f"USER: <image>\n{image_prompt} ASSISTANT:" | |
| # Assuming your model can handle PIL images | |
| image = image.convert("RGB") # Ensure image is RGB mode | |
| inputs = processor(text=prompt, images=image, return_tensors="pt") | |
| generate_ids = base_model.generate(**inputs, max_new_tokens=1024) | |
| decoded_response = processor.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] | |
| return decoded_response | |
| def image_to_base64(image_path): | |
| with open(image_path, 'rb') as img: | |
| encoded_string = base64.b64encode(img.read()) | |
| return encoded_string.decode('utf-8') | |
| # Function that takes User Inputs and displays it on ChatUI | |
| def query_message(history,txt,img): | |
| image_prompt = deep_translator_bn_en(txt) | |
| history += [(image_prompt,None)] | |
| base64 = image_to_base64(img) | |
| data_url = f"data:image/jpeg;base64,{base64}" | |
| history += [(f"{image_prompt} ", None)] | |
| return history | |
| # Function that takes User Inputs, generates Response and displays on Chat UI | |
| def llm_response(history,text,img): | |
| image_prompt = deep_translator_bn_en(text) | |
| response = inference(img,image_prompt) | |
| assistant_index = response.find("ASSISTANT:") | |
| extracted_string = response[assistant_index + len("ASSISTANT:"):].strip() | |
| output = deep_translator_en_bn(extracted_string) | |
| history += [(text,output)] | |
| return history | |
| # Interface Code | |
| st.title('My_BoT') | |
| # Create a sidebar | |
| sidebar = st.sidebar | |
| sidebar.header('User Inputs') | |
| # Add a file uploader to the sidebar | |
| uploaded_file = sidebar.file_uploader("Upload an Image", type=['png', 'jpg', 'jpeg']) | |
| # Add a text input to the sidebar | |
| text_input = sidebar.text_input("Enter text and press enter") | |
| # Initialize session state for history if it doesn't exist | |
| if 'history' not in st.session_state: | |
| st.session_state.history = [] | |
| # Check if text is entered and no image is uploaded | |
| if text_input and uploaded_file is None: | |
| st.write("Please upload an image.") | |
| # Add a button to the sidebar | |
| submit_button = sidebar.button("Submit") | |
| # When the button is clicked, generate the response and display it | |
| if submit_button: | |
| if uploaded_file is not None: | |
| image = PIL.Image.open(uploaded_file) | |
| st.session_state.history = llm_response(st.session_state.history, text_input, image) | |
| for text, output in st.session_state.history: | |
| st.write(f"User: {text}") | |
| if output is not None: | |
| st.write(f"Assistant: {output}") |