Update app.py

app.py

@@ -8,8 +8,8 @@ import transformers  # LLM
 PUBLICATIONS_TO_RETRIEVE = 5
 
 # The template for the RAG prompt
-RAG_TEMPLATE = """You are an AI assistant who enjoys helping users learn about research. 
-Answer the USER_QUERY on additive manufacturing research using the RESEARCH_EXCERPTS. 
+RAG_TEMPLATE = """You are an AI assistant who enjoys helping users learn about research.
+Answer the USER_QUERY on additive manufacturing research using the RESEARCH_EXCERPTS.
 Provide a concise ANSWER based on these excerpts. Avoid listing references.
 ===== RESEARCH_EXCERPTS =====
 {research_excerpts}
@@ -18,6 +18,23 @@ Provide a concise ANSWER based on these excerpts. Avoid listing references.
 ===== ANSWER =====
 """
 
+# Example Queries for Interface
+EXAMPLE_QUERIES = [
+    {"text": "What is multi-material 3D printing?"},
+    {"text": "How is additive manufacturing being applied in aerospace?"},
+    {"text": "Tell me about innovations in metal 3D printing techniques."},
+    {"text": "What are some sustainable materials for 3D printing?"},
+    {"text": "What are the challenges with support structures in 3D printing?"},
+    {"text": "How is 3D printing impacting the medical field?"},
+    {"text": "What are common applications of additive manufacturing in industry?"},
+    {"text": "What are the benefits and limitations of using polymers in 3D printing?"},
+    {"text": "Tell me about the environmental impacts of additive manufacturing."},
+    {"text": "What are the primary limitations of current 3D printing technologies?"},
+    {"text": "How are researchers improving the speed of 3D printing processes?"},
+    {"text": "What are best practices for post-processing in additive manufacturing?"},
+]
+
+
 # Load vectorstore of SFF publications
 publication_vectorstore = langchain_community.vectorstores.FAISS.load_local(
     folder_path="publication_vectorstore",
@@ -30,13 +47,17 @@ publication_vectorstore = langchain_community.vectorstores.FAISS.load_local(
 )
 
 # Create the callable LLM
+model = transformers.AutoModelForCausalLM.from_pretrained(
+    "Qwen/Qwen2.5-7B-Instruct-AWQ"
+)
+model.to("cuda")  # Move the model to GPU
+tokenizer = transformers.AutoTokenizer.from_pretrained("Qwen/Qwen2.5-7B-Instruct-AWQ")
+
 llm = transformers.pipeline(
     task="text-generation",
-    model="Qwen/Qwen2.5-7B-Instruct-AWQ",
-    device=0,
-    streamer=transformers.TextStreamer(
-        transformers.AutoTokenizer.from_pretrained("Qwen/Qwen2.5-7B-Instruct-AWQ")
-    ),
+    model=model,
+    tokenizer=tokenizer,
+    device=0,  # Ensure the model is loaded on the GPU
 )
 
 
@@ -90,22 +111,6 @@ def reply(message: str, history: list[str]) -> str:
     return response[0]["generated_text"].strip("= ")
 
 
-# Example Queries for Interface
-EXAMPLE_QUERIES = [
-    {"text": "What is multi-material 3D printing?"},
-    {"text": "How is additive manufacturing being applied in aerospace?"},
-    {"text": "Tell me about innovations in metal 3D printing techniques."},
-    {"text": "What are some sustainable materials for 3D printing?"},
-    {"text": "What are the challenges with support structures in 3D printing?"},
-    {"text": "How is 3D printing impacting the medical field?"},
-    {"text": "What are common applications of additive manufacturing in industry?"},
-    {"text": "What are the benefits and limitations of using polymers in 3D printing?"},
-    {"text": "Tell me about the environmental impacts of additive manufacturing."},
-    {"text": "What are the primary limitations of current 3D printing technologies?"},
-    {"text": "How are researchers improving the speed of 3D printing processes?"},
-    {"text": "What are best practices for post-processing in additive manufacturing?"},
-]
-
 # Run the Gradio Interface
 gradio.ChatInterface(
     reply,