Spaces:
Sleeping
Sleeping
viktor-hirenko commited on
Commit ·
5ab05da
1
Parent(s): a601d1a
fix: handle StopIteration in HF Inference API streaming
Browse files- llm_handler.py +14 -3
llm_handler.py
CHANGED
|
@@ -112,19 +112,30 @@ class LLMHandler:
|
|
| 112 |
stream = self.client.text_generation(
|
| 113 |
prompt=full_prompt,
|
| 114 |
model=self.model,
|
| 115 |
-
max_new_tokens=
|
| 116 |
temperature=0.7,
|
| 117 |
top_p=0.95,
|
| 118 |
stream=True,
|
| 119 |
details=False,
|
| 120 |
)
|
| 121 |
|
|
|
|
|
|
|
| 122 |
for token in stream:
|
| 123 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 124 |
|
|
|
|
|
|
|
|
|
|
| 125 |
except Exception as e:
|
| 126 |
logger.error(f"Error streaming answer: {e}")
|
| 127 |
-
|
| 128 |
|
| 129 |
|
| 130 |
def format_response(
|
|
|
|
| 112 |
stream = self.client.text_generation(
|
| 113 |
prompt=full_prompt,
|
| 114 |
model=self.model,
|
| 115 |
+
max_new_tokens=512,
|
| 116 |
temperature=0.7,
|
| 117 |
top_p=0.95,
|
| 118 |
stream=True,
|
| 119 |
details=False,
|
| 120 |
)
|
| 121 |
|
| 122 |
+
# Collect tokens and yield them
|
| 123 |
+
has_content = False
|
| 124 |
for token in stream:
|
| 125 |
+
if token: # Only yield non-empty tokens
|
| 126 |
+
has_content = True
|
| 127 |
+
yield token
|
| 128 |
+
|
| 129 |
+
# If no content was generated, yield a fallback message
|
| 130 |
+
if not has_content:
|
| 131 |
+
yield "I apologize, but I couldn't generate an answer based on the provided context. Please try rephrasing your question."
|
| 132 |
|
| 133 |
+
except StopIteration:
|
| 134 |
+
# Handle empty generator
|
| 135 |
+
yield "I apologize, but I couldn't generate an answer. The model returned an empty response."
|
| 136 |
except Exception as e:
|
| 137 |
logger.error(f"Error streaming answer: {e}")
|
| 138 |
+
yield f"\n\n❌ Error: {str(e)}"
|
| 139 |
|
| 140 |
|
| 141 |
def format_response(
|