Daniel Marques commited on
Commit
2a13ed4
1 Parent(s): fac2b5c

feat: add backend

Browse files
Files changed (1) hide show
  1. main.py +4 -14
main.py CHANGED
@@ -2,8 +2,6 @@ import os
2
  import glob
3
  import shutil
4
  import subprocess
5
- import contextvars
6
- import asyncio
7
 
8
  from typing import Any, Dict, List
9
 
@@ -17,7 +15,7 @@ from langchain.chains import RetrievalQA
17
  from langchain.embeddings import HuggingFaceInstructEmbeddings
18
  from langchain.prompts import PromptTemplate
19
  from langchain.memory import ConversationBufferMemory
20
- from langchain.callbacks.base import BaseCallbackHandler, AsyncCallbackHandler
21
  from langchain.schema import LLMResult
22
 
23
  # from langchain.embeddings import HuggingFaceEmbeddings
@@ -34,15 +32,8 @@ class Predict(BaseModel):
34
  class Delete(BaseModel):
35
  filename: str
36
 
37
- websocket_state = None
38
-
39
- class MyCustomSyncHandler(AsyncCallbackHandler):
40
  def on_llm_new_token(self, token: str, **kwargs) -> None:
41
- print(f"{websocket_state}")
42
-
43
- asyncio.sleep(1.5)
44
- websocket_state.send_text(f"token: {token}")
45
-
46
  print(f"token: {token}")
47
 
48
  # if torch.backends.mps.is_available():
@@ -76,13 +67,13 @@ Always answer in the most helpful and safe way possible.
76
  If you don't know the answer to a question, just say that you don't know, don't try to make up an answer, don't share false information.
77
  Use 15 sentences maximum. Keep the answer as concise as possible.
78
  Always say "thanks for asking!" at the end of the answer.
79
- Context: {history} \n {context}
80
  Question: {question}
81
  """
82
 
83
  memory = ConversationBufferMemory(input_key="question", memory_key="history")
84
 
85
- QA_CHAIN_PROMPT = PromptTemplate(input_variables=["history", "context", "question"], template=template)
86
 
87
  QA = RetrievalQA.from_chain_type(
88
  llm=LLM,
@@ -91,7 +82,6 @@ QA = RetrievalQA.from_chain_type(
91
  return_source_documents=SHOW_SOURCES,
92
  chain_type_kwargs={
93
  "prompt": QA_CHAIN_PROMPT,
94
- "memory": memory
95
  },
96
  )
97
 
 
2
  import glob
3
  import shutil
4
  import subprocess
 
 
5
 
6
  from typing import Any, Dict, List
7
 
 
15
  from langchain.embeddings import HuggingFaceInstructEmbeddings
16
  from langchain.prompts import PromptTemplate
17
  from langchain.memory import ConversationBufferMemory
18
+ from langchain.callbacks.base import BaseCallbackHandler
19
  from langchain.schema import LLMResult
20
 
21
  # from langchain.embeddings import HuggingFaceEmbeddings
 
32
  class Delete(BaseModel):
33
  filename: str
34
 
35
+ class MyCustomSyncHandler(BaseCallbackHandler):
 
 
36
  def on_llm_new_token(self, token: str, **kwargs) -> None:
 
 
 
 
 
37
  print(f"token: {token}")
38
 
39
  # if torch.backends.mps.is_available():
 
67
  If you don't know the answer to a question, just say that you don't know, don't try to make up an answer, don't share false information.
68
  Use 15 sentences maximum. Keep the answer as concise as possible.
69
  Always say "thanks for asking!" at the end of the answer.
70
+ Context: {context}
71
  Question: {question}
72
  """
73
 
74
  memory = ConversationBufferMemory(input_key="question", memory_key="history")
75
 
76
+ QA_CHAIN_PROMPT = PromptTemplate(input_variables=["context", "question"], template=template)
77
 
78
  QA = RetrievalQA.from_chain_type(
79
  llm=LLM,
 
82
  return_source_documents=SHOW_SOURCES,
83
  chain_type_kwargs={
84
  "prompt": QA_CHAIN_PROMPT,
 
85
  },
86
  )
87