Update app.py

app.py

@@ -1,4 +1,4 @@
-import time, aiohttp, asyncio, json, os, multiprocessing, torch, spaces
+import time, aiohttp, asyncio, json, os, multiprocessing, torch
 from minivectordb.embedding_model import EmbeddingModel
 from minivectordb.vector_database import VectorDatabase
 from text_util_en_pt.cleaner import structurize_text, detect_language, Language
@@ -12,10 +12,18 @@ import gradio as gr
 torch.set_num_threads(2)
 
 model = EmbeddingModel(use_quantized_onnx_model=True)
-tokenizer = AutoTokenizer.from_pretrained("GeneZC/MiniChat-3B", use_fast=False)
-llm = AutoModelForCausalLM.from_pretrained("GeneZC/MiniChat-3B")
+tokenizer = AutoTokenizer.from_pretrained("sreeramajay/TinyLlama-1.1B-orca-v1.0")
+llm = AutoModelForCausalLM.from_pretrained("sreeramajay/TinyLlama-1.1B-orca-v1.0")
 
-prompt_template = "<s> [|User|] $PROMPT </s>[|Assistant|]"
+prompt_template = """<|system|>
+You are a helpful assistant chatbot.</s>
+<|user|>
+$PROMPT</s>
+<|assistant|>
+"""
+
+def count_tokens(text):
+    return len(tokenizer.encode(text))
 
 def fetch_links(query, max_results=5):
     with DDGS() as ddgs:
@@ -59,7 +67,6 @@ def retrieval_pipeline(query):
 
     return context, websearch_time, webcrawl_time, embedding_time, retrieval_time, links
 
-@spaces.GPU(enable_queue=True)
 def ask_open_llm(prompt):
     device = torch.device('cuda')
 
@@ -73,7 +80,7 @@ def ask_open_llm(prompt):
     generate_kwargs = dict(
         model_inputs,
         streamer=streamer,
-        max_new_tokens=512,
+        max_new_tokens=2048 - count_tokens(prompt),
         top_p=0.2,
         top_k=20,
         temperature=0.4,