Update app.py

app.py

@@ -1,52 +1,12 @@
-import streamlit as st
-from transformers import AutoTokenizer, AutoModelForCausalLM
-import PyPDF2
-
-# Model Identifier
-model_name = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(model_name)
-
-st.set_page_config(page_title="DOCAI with TinyLlama", page_icon="🤖", layout="wide")
-st.title("📄 Document Chat with TinyLlama 🤖")
-
-# Initialize chat history
-if "messages" not in st.session_state:
-    st.session_state["messages"] = [{"role": "assistant", "content": "How can I help you with your document?"}]
-
-def process_document_and_respond(uploaded_file, question):
-    if uploaded_file:
-        try:
-            reader = PyPDF2.PdfReader(uploaded_file)
-            text_content = ""
-            for page in reader.pages:
-                text_content += page.extract_text()
-
-            # Basic context injection - you might need a more sophisticated approach
-            prompt = f"Context: {text_content}\n\nQuestion: {question}\n\nAnswer:"
-            inputs = tokenizer([prompt], return_tensors="pt").to(model.device)
-            outputs = model.generate(**inputs, max_new_tokens=512, do_sample=True, top_p=0.9, top_k=50, temperature=0.7)
-            response = tokenizer.batch_decode(outputs, skip_special_tokens=True)[0]
-            return response.split("Answer:")[-1].strip()
-
-        except Exception as e:
-            return f"Error processing document: {e}"
-    else:
-        return "Please upload a PDF document first."
-
-# Upload the PDF
-uploaded_file = st.file_uploader("Upload a PDF document", type="pdf")
-
-# Input for the question
-question = st.text_input("Ask your question about the document:")
-
-if question:
-    with st.spinner("Thinking..."):
-        response = process_document_and_respond(uploaded_file, question)
-        st.session_state["messages"].append({"role": "user", "content": question})
-        st.session_state["messages"].append({"role": "assistant", "content": response})
-
-# Display chat history
-for message in st.session_state["messages"]:
-    with st.chat_message(message["role"]):
-        st.markdown(message["content"])
+import os
+from langchain_google_genai import ChatGoogleGenerativeAI
+from dotenv import load_dotenv
+
+load_dotenv()
+llm = ChatGoogleGenerativeAI(
+    model="models/gemini-pro",
+    temperature=0.7
+)
+
+response = llm.invoke("Hello, what is AI?")
+print(response.content)