Update app.py

app.py

@@ -1,37 +1,3 @@
-"""
-This script sets up a Gradio interface for querying an AI assistant about additive manufacturing research.
-It uses a vectorstore to retrieve relevant research excerpts and a language model to generate responses.
-
-Modules:
-    - gradio: Interface handling
-    - spaces: For GPU
-    - transformers: LLM Loading
-    - langchain_community.vectorstores: Vectorstore for publications
-    - langchain_huggingface: Embeddings
-
-Constants:
-    - PUBLICATIONS_TO_RETRIEVE: The number of publications to retrieve for the prompt
-    - RAG_TEMPLATE: The template for the RAG prompt
-
-Functions:
-    - preprocess(query: str) -> str: Generates a prompt based on the top k documents matching the query.
-    - reply(message: str, history: list[str]) -> str: Generates a response to the user’s message.
-
-Example Queries:
-    - "What is multi-material 3D printing?"
-    - "How is additive manufacturing being applied in aerospace?"
-    - "Tell me about innovations in metal 3D printing techniques."
-    - "What are some sustainable materials for 3D printing?"
-    - "What are the biggest challenges with support structures in additive manufacturing?"
-    - "How is 3D printing impacting the medical field?"
-    - "What are some common applications of additive manufacturing in industry?"
-    - "What are the benefits and limitations of using polymers in 3D printing?"
-    - "Tell me about the environmental impacts of additive manufacturing."
-    - "What are the primary limitations of current 3D printing technologies?"
-    - "How are researchers improving the speed of 3D printing processes?"
-    - "What are the best practices for managing post-processing in additive manufacturing?"
-"""
-
 import gradio  # Interface handling
 import spaces  # For GPU
 import langchain_community.vectorstores  # Vectorstore for publications
@@ -65,14 +31,13 @@ publication_vectorstore = langchain_community.vectorstores.FAISS.load_local(
     ),
     allow_dangerous_deserialization=True,
 )
-
-# Create the callable LLM
-llm = transformers.pipeline(
-    task="text-generation", 
-    model="Qwen/Qwen2.5-7B-Instruct-AWQ", 
-    device="cuda", 
-    streamer=transformers.TextStreamer(transformers.AutoTokenizer.from_pretrained("Qwen/Qwen2.5-7B-Instruct-AWQ"))
-)
+#
+# # Create the callable LLM
+# llm = transformers.pipeline(
+#     task="text-generation",
+#     model="Qwen/Qwen2.5-7B-Instruct-AWQ",
+#     device="cuda",
+# )
 
 
 def preprocess(query: str) -> str:
@@ -99,12 +64,12 @@ def preprocess(query: str) -> str:
         research_excerpts="\n\n".join(research_excerpts), query=query
     )
 
-    # Print the prompt for debugging purposes
-    print(prompt)
-
     return prompt
 
 
+import threading
+
+
 @spaces.GPU
 def reply(message: str, history: list[str]) -> str:
     """
@@ -118,13 +83,30 @@ def reply(message: str, history: list[str]) -> str:
         str: The generated response from the language model.
     """
 
-    yield llm(
-        preprocess(message),
-        max_new_tokens=512,
-        return_full_text=False,
-    )[
-        0
-    ]["generated_text"]
+    tok = transformers.AutoTokenizer.from_pretrained("Qwen/Qwen2.5-7B-Instruct-AWQ")
+    model = transformers.AutoModelForCausalLM.from_pretrained(
+        "Qwen/Qwen2.5-7B-Instruct-AWQ"
+    )
+    inputs = tok([preprocess(message)], return_tensors="pt")
+    streamer = transformers.TextIteratorStreamer(tok)
+
+    generation_kwargs = dict(inputs, streamer=streamer, max_new_tokens=512)
+    thread = threading.Thread(target=model.generate, kwargs=generation_kwargs)
+    thread.start()
+    generated_text = ""
+    for new_text in streamer:
+        generated_text += new_text
+        yield generated_text
+
+    # yield llm(
+    #     preprocess(message),
+    #     max_new_tokens=512,
+    #     return_full_text=False,
+    #     streamer=transformers.TextIteratorStreamer(
+    #         transformers.AutoTokenizer.from_pretrained("Qwen/Qwen2.5-7B-Instruct-AWQ")
+    #     ),
+    # )[0]["generated_text"]
+
 
 # Example Queries for Interface
 EXAMPLE_QUERIES = [
@@ -132,14 +114,20 @@ EXAMPLE_QUERIES = [
     {"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 biggest challenges with support structures in additive manufacturing?"},
+    {
+        "text": "What are the biggest challenges with support structures in additive manufacturing?"
+    },
     {"text": "How is 3D printing impacting the medical field?"},
-    {"text": "What are some common applications of additive manufacturing in industry?"},
+    {
+        "text": "What are some 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 the best practices for managing post-processing in additive manufacturing?"}
+    {
+        "text": "What are the best practices for managing post-processing in additive manufacturing?"
+    },
 ]
 
 # Run the Gradio Interface