model download issues

model/model_runner.py

@@ -35,7 +35,7 @@ def load_model():
 
 def generate(prompt: str, max_tokens: int = 256) -> str:
     model = load_model()
-    response = model(
+    output = model(
         prompt,
         max_tokens=max_tokens,
         stop=["\nUser:", "\nAssistant:"],
@@ -43,4 +43,13 @@ def generate(prompt: str, max_tokens: int = 256) -> str:
         temperature=0.7,
         top_p=0.95,
     )
+    raw_text = output["choices"][0]["text"]
+
+    # Truncate at first stop sequence manually
+    for stop_token in ["\nUser:", "\nAssistant:"]:
+        if stop_token in raw_text:
+            raw_text = raw_text.split(stop_token)[0]
+
+    return raw_text.strip()
+
     return response["choices"][0]["text"].strip()