Update app.py

app.py

@@ -336,8 +336,8 @@ def answer_query(query_text, index, documents, llm_model, llm_tokenizer, embeddi
     retrieved_info = get_retrieved_info(documents, I, D)
     formatted_info = format_retrieved_info(retrieved_info)
     prompt = generate_prompt(query_text, formatted_info)
-    answer = answer_using_gemma(prompt, llm_model, llm_tokenizer)
-    return answer
+    # answer = answer_using_gemma(prompt, llm_model, llm_tokenizer)
+    return prompt
 
 
 
@@ -393,7 +393,30 @@ model = AutoModelForCausalLM.from_pretrained("google/gemma-2b-it", token=HF_TOKE
 
 
 def make_inference(query, hist):
-    return answer_query(query, index, documents, model, tokenizer, CFG.embedding_model, CFG.n_samples, CFG.device)
+    prompt = answer_query(query, index, documents, model, tokenizer, CFG.embedding_model, CFG.n_samples, CFG.device)
+    # answer = answer_using_gemma(prompt, llm_model, llm_tokenizer)
+    model_inputs = tokenizer(prompt, return_tensors="pt").to("cuda" if torch.cuda.is_available() else "cpu")
+    count_tokens = lambda text: len(tokenizer.tokenize(text))
+    
+    streamer = TextIteratorStreamer(tokenizer, timeout=540., skip_prompt=True, skip_special_tokens=True)
+    
+    generate_kwargs = dict(
+        model_inputs,
+        streamer=streamer,
+        max_new_tokens=6000 - count_tokens(prompt),
+        top_p=0.2,
+        top_k=20,
+        temperature=0.1,
+        repetition_penalty=2.0,
+        length_penalty=-0.5,
+        num_beams=1
+    )
+    t = Thread(target=model.generate, kwargs=generate_kwargs)
+    t.start()  # Starting the generation in a separate thread.
+    partial_message = ""
+    for new_token in streamer:
+        partial_message += new_token
+        yield partial_message
 
 demo = gr.ChatInterface(fn = make_inference, 
                 examples = ["What is diabetes?", "Is ginseng good for diabetes?", "What are the symptoms of diabetes?", "What is Celiac disease?"],