vilarin commited on
Commit
85c845b
·
verified ·
1 Parent(s): 42681ce

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +6 -12
app.py CHANGED
@@ -37,8 +37,8 @@ def terminate():
37
 
38
  import ollama
39
  import gradio as gr
40
- from ollama import Client
41
- client = Client(host='http://localhost:11434', timeout=120)
42
 
43
  HF_TOKEN = os.environ.get("HF_TOKEN", None)
44
 
@@ -115,7 +115,7 @@ async def stream_chat(message: str, history: list, model: str, temperature: floa
115
 
116
  print(f"Conversation is -\n{conversation}")
117
 
118
- response = client.chat(
119
  model=model,
120
  stream=True,
121
  messages=conversation,
@@ -128,14 +128,8 @@ async def stream_chat(message: str, history: list, model: str, temperature: floa
128
  'repeat_penalty': penalty,
129
  'low_vram': True,
130
  },
131
- )
132
-
133
- print(response)
134
-
135
- buffer = ""
136
- for chunk in response:
137
- buffer += chunk["message"]["content"]
138
- yield buffer
139
 
140
 
141
  async def main(message: str, history: list, model: str, temperature: float, max_new_tokens: int, top_p: float, top_k: int, penalty: float):
@@ -145,7 +139,7 @@ async def main(message: str, history: list, model: str, temperature: float, max_
145
  else:
146
  if not INIT_SIGN:
147
  yield "Please initialize Ollama"
148
- else:
149
  async for response in stream_chat(
150
  message,
151
  history,
 
37
 
38
  import ollama
39
  import gradio as gr
40
+ from ollama import AsyncClient
41
+ client = AsyncClient(host='http://localhost:11434', timeout=120)
42
 
43
  HF_TOKEN = os.environ.get("HF_TOKEN", None)
44
 
 
115
 
116
  print(f"Conversation is -\n{conversation}")
117
 
118
+ async for part in await client.chat(
119
  model=model,
120
  stream=True,
121
  messages=conversation,
 
128
  'repeat_penalty': penalty,
129
  'low_vram': True,
130
  },
131
+ ):
132
+ yield part['message']['content']
 
 
 
 
 
 
133
 
134
 
135
  async def main(message: str, history: list, model: str, temperature: float, max_new_tokens: int, top_p: float, top_k: int, penalty: float):
 
139
  else:
140
  if not INIT_SIGN:
141
  yield "Please initialize Ollama"
142
+ else:
143
  async for response in stream_chat(
144
  message,
145
  history,