Spaces:
Paused
Paused
fix use after free
Browse files
app.py
CHANGED
|
@@ -102,12 +102,15 @@ def completion(prompt: str, model_id: str):
|
|
| 102 |
temperature=None,
|
| 103 |
)
|
| 104 |
|
|
|
|
|
|
|
|
|
|
| 105 |
# cleanup memory
|
| 106 |
del model, tokenizer
|
| 107 |
torch.cuda.empty_cache()
|
| 108 |
gc.collect()
|
| 109 |
|
| 110 |
-
return
|
| 111 |
|
| 112 |
|
| 113 |
def completion_openrouter(prompt: str, model_id: str):
|
|
|
|
| 102 |
temperature=None,
|
| 103 |
)
|
| 104 |
|
| 105 |
+
# decode response
|
| 106 |
+
resp = tokenizer.decode(outputs[0][input_ids.shape[-1]:], skip_special_tokens=True)
|
| 107 |
+
|
| 108 |
# cleanup memory
|
| 109 |
del model, tokenizer
|
| 110 |
torch.cuda.empty_cache()
|
| 111 |
gc.collect()
|
| 112 |
|
| 113 |
+
return resp
|
| 114 |
|
| 115 |
|
| 116 |
def completion_openrouter(prompt: str, model_id: str):
|