Update app.py

app.py

@@ -1,6 +1,8 @@
 import gradio as gr
 import numpy as np
 from sklearn.metrics.pairwise import cosine_similarity
+import subprocess
+
 # from sklearn.decomposition import PCA
 from langchain_community.llms import Ollama
 from langchain_chroma import Chroma
@@ -10,26 +12,38 @@ from langchain_community.document_loaders import DirectoryLoader, TextLoader, Py
 from langchain_experimental.text_splitter import SemanticChunker
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain_community.embeddings.ollama import OllamaEmbeddings
+from langchain.embeddings import HuggingFaceEmbeddings
 
 from typing import List, Dict
 from langchain.docstore.document import Document
 
-from transformers import T5Tokenizer, T5ForConditionalGeneration
-
+from transformers import T5Tokenizer, T5ForConditionalGeneration, AutoTokenizer, AutoModelForCausalLM, pipeline
 tokenizer = T5Tokenizer.from_pretrained("Voicelab/vlt5-base-keywords")
 model = T5ForConditionalGeneration.from_pretrained("Voicelab/vlt5-base-keywords")
 
+model_name = "meta-llama/Llama-3.2-1B-Instruct"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto")
+
+generator = pipeline(
+    "text-generation",
+    model=model,
+    tokenizer=tokenizer,
+    max_new_tokens=500,  # Adjust as needed
+    temperature=0.5      # Adjust as needed
+)
+
 vectorstore = Chroma(
     # docs,
-    embedding_function=OllamaEmbeddings(model = "gemma:2b"),
+    embedding_function=HuggingFaceEmbeddings(model_name = "google/gemma-2b"),
     persist_directory="chroma_db"
 )
 
 print(vectorstore.similarity_search_with_score("Course Leader"))
 
-llm = Ollama(
-    model="llama3.2:3b"
-)
+# llm = Ollama(
+#     model="llama3.2:3b"
+# )
 
 def retrieve_relevant_chunks(
     vector_store: Chroma,
@@ -121,7 +135,7 @@ def main(query):
         print(f"Error: {str(e)}")
 
     formatted_prompt = f"""
-        You are an AI assistant. Your goal is to answer questions regarding student handbooks based on the following context provided. Make sure all the answers are within the given context:
+        You are an AI assistant. Your goal is to answer questions regarding degree information based on the following context provided. Make sure all the answers are within the given context and act like you are a representative of IIT so do not mention anthing for users to know that you are reading something:
         {prompt}
 
         Based on the above, answer the following question:
@@ -129,7 +143,7 @@ def main(query):
         Give the answer in a clear and concise manner
         """
 
-    response = llm.predict(formatted_prompt)
+    response = generator(formatted_prompt, return_full_text=False)
 
     return response
 
@@ -153,5 +167,4 @@ with gr.Blocks() as demo:
     login_btn.click(main, inputs=[query], outputs=answer)
    
 # demo.launch(share = True, auth=authenticate)
-demo.launch(share = True)
-
+demo.launch()