VisionLang / backend.py
parthib07's picture
Upload 6 files
2098354 verified
import google.generativeai as genai
from llama_index.llms.gemini import Gemini
from llama_index.embeddings.gemini import GeminiEmbedding
import os
import tempfile
from llama_index.core import VectorStoreIndex, SimpleDirectoryReader
from llama_index.core import Settings
import time
from google.api_core.exceptions import GoogleAPIError
import streamlit as st
genai.configure(api_key=os.environ.get("GOOGLE_API_KEY"))
llm = Gemini(model_name="models/gemini-1.5-pro")
embeddings = GeminiEmbedding(model_name="models/embedding-001")
def normal_response(query):
prompt = """You are a helpful Bot named VisionLang Build by Parthib Karak.
Given a question, generate answer based on the Question.
Question: {question}
"""
try:
response = llm.complete(prompt + query)
return response.text
except GoogleAPIError as e:
return f"Error generating response: {str(e)}"
def uploaded_file_to_response(file, query):
file_extension = os.path.splitext(file.name)[-1].lower()
try:
if file_extension in [".pdf", ".docx", ".txt", ".py", ".js", ".java", ".cpp"]:
temp_dir = tempfile.mkdtemp()
temp_file_path = os.path.join(temp_dir, file.name)
with open(temp_file_path, "wb") as f:
f.write(file.read())
document = SimpleDirectoryReader(temp_dir)
data = document.load_data()
Settings.llm = llm
Settings.embed_model = embeddings
index = VectorStoreIndex.from_documents(data, settings=Settings)
query_engine = index.as_query_engine()
response = query_engine.query(query)
return response
elif file_extension in [".mp4", ".avi", ".mov",".mkv"]:
temp_dir = tempfile.mkdtemp()
temp_file_path = os.path.join(temp_dir, file.name)
with open(temp_file_path, "wb") as f:
f.write(file.read())
uploaded_file = genai.upload_file(temp_file_path, mime_type="video/mp4")
st.success("video uploaded successfully")
time.sleep(2)
response = llm.complete([query, uploaded_file])
return response.text
elif file_extension in [".png", ".jpg", ".jpeg"]:
uploaded_file = genai.upload_file(file, mime_type="image/jpeg")
time.sleep(2)
response = llm.complete([query, uploaded_file])
return response.text
else:
uploaded_file = genai.upload_file(file, mime_type="application/octet-stream")
time.sleep(2)
response = llm.complete([query, uploaded_file])
return response.text
except GoogleAPIError as e:
return f"Error processing file: {str(e)}"