Update app.py

app.py

@@ -2,28 +2,14 @@ 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
-from transformers import AutoModelForCausalLM, AutoTokenizer
-from transformers import StoppingCriteria, StoppingCriteriaList, TextIteratorStreamer
 from webtextcrawler.webtextcrawler import extract_text_from_url
-from threading import Thread
 from duckduckgo_search import DDGS
 import gradio as gr
 
 torch.set_num_threads(2)
 
+openrouter_key = os.environ.get("OPENROUTER_KEY")
 model = EmbeddingModel(use_quantized_onnx_model=True)
-tokenizer = AutoTokenizer.from_pretrained("sreeramajay/TinyLlama-1.1B-orca-v1.0")
-llm = AutoModelForCausalLM.from_pretrained("sreeramajay/TinyLlama-1.1B-orca-v1.0")
-
-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:
@@ -50,7 +36,7 @@ def index_and_search(query, text):
 
     # Retrieval
     start = time.time()
-    search_results = vector_db.find_most_similar(query_embedding, k = 5)
+    search_results = vector_db.find_most_similar(query_embedding, k = 12)
     retrieval_time = time.time() - start
     return '\n'.join([s['sentence'] for s in search_results[2]]), embedding_time, retrieval_time
 
@@ -67,30 +53,7 @@ def retrieval_pipeline(query):
 
     return context, websearch_time, webcrawl_time, embedding_time, retrieval_time, links
 
-def ask_open_llm(prompt):
-    model_inputs = tokenizer([
-        prompt
-    ], return_tensors="pt")
-    
-    streamer = TextIteratorStreamer(tokenizer, timeout=120., skip_prompt=True, skip_special_tokens=True)
-    
-    generate_kwargs = dict(
-        model_inputs,
-        streamer=streamer,
-        max_new_tokens=2048 - count_tokens(prompt),
-        do_sample=True,
-        temperature=0.7,
-        top_p=0.9,
-        repetition_penalty=2.5
-    )
-    t = Thread(target=llm.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
-
-def predict(message, history):
+async def predict(message, history):
     context, websearch_time, webcrawl_time, embedding_time, retrieval_time, links = retrieval_pipeline(message)
 
     if detect_language(message) == Language.ptbr:
@@ -98,14 +61,49 @@ def predict(message, history):
     else:
         prompt = f"Context:\n\n{context}\n\nBased on the context, answer: {message}"
 
-    prompt = prompt_template.replace("$PROMPT", prompt)
+    url = "https://openrouter.ai/api/v1/chat/completions"
+    headers = { "Content-Type": "application/json",
+                "Authorization": f"Bearer {openrouter_key}" }
+    body = { "stream": True,
+             "models": [
+                    "mistralai/mistral-7b-instruct:free", 
+                    "nousresearch/nous-capybara-7b:free",
+                    "huggingfaceh4/zephyr-7b-beta:free"
+             ],
+             "route": "fallback",
+             "max_tokens": 768,
+             "messages": [
+                 {"role": "user", "content": prompt}
+             ] }
 
     full_response = ""
-    for partial_message in ask_open_llm(prompt):
-        full_response += partial_message
-        yield full_response
+    async with aiohttp.ClientSession() as session:
+        async with session.post(url, headers=headers, json=body) as response:
+            buffer = ""  # A buffer to hold incomplete lines of data
+            async for chunk in response.content.iter_any():
+                buffer += chunk.decode()
+                while "\n" in buffer:  # Process as long as there are complete lines in the buffer
+                    line, buffer = buffer.split("\n", 1)
+
+                    if line.startswith("data: "):
+                        event_data = line[len("data: "):]
+                        if event_data != '[DONE]':
+                            try:
+                                current_text = json.loads(event_data)['choices'][0]['delta']['content']
+                                full_response += current_text
+                                yield full_response
+                                await asyncio.sleep(0.01)
+                            except Exception:
+                                try:
+                                    current_text = json.loads(event_data)['choices'][0]['text']
+                                    full_response += current_text
+                                    yield full_response
+                                    await asyncio.sleep(0.01)
+                                except Exception:
+                                    pass
     
     final_metadata_block = ""
+
     final_metadata_block += f"Links visited:\n"
     for link in links:
         final_metadata_block += f"{link}\n"