Spaces:
Sleeping
Sleeping
no message
Browse files
main.py
CHANGED
@@ -22,32 +22,43 @@ app = FastAPI()
|
|
22 |
# Initialize the InferenceClient with your model
|
23 |
client = InferenceClient("mistralai/Mixtral-8x7B-Instruct-v0.1")
|
24 |
|
|
|
25 |
class Item(BaseModel):
|
26 |
prompt: str
|
27 |
history: list
|
28 |
system_prompt: str
|
29 |
temperature: float = 0.8
|
30 |
-
max_new_tokens: int =
|
31 |
top_p: float = 0.15
|
32 |
repetition_penalty: float = 1.0
|
33 |
|
34 |
-
def format_prompt(
|
35 |
-
|
36 |
-
for
|
37 |
-
|
38 |
-
|
39 |
-
|
40 |
-
|
|
|
|
|
|
|
41 |
|
42 |
def generate_stream(item: Item) -> Generator[bytes, None, None]:
|
43 |
formatted_prompt = format_prompt(f"{item.system_prompt}, {item.prompt}", item.history)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
44 |
generate_kwargs = {
|
45 |
"temperature": item.temperature,
|
46 |
-
"max_new_tokens":
|
47 |
"top_p": item.top_p,
|
48 |
"repetition_penalty": item.repetition_penalty,
|
49 |
"do_sample": True,
|
50 |
-
"seed": 42,
|
51 |
}
|
52 |
|
53 |
# Stream the response from the InferenceClient
|
@@ -59,11 +70,17 @@ def generate_stream(item: Item) -> Generator[bytes, None, None]:
|
|
59 |
}
|
60 |
yield json.dumps(chunk).encode("utf-8") + b"\n"
|
61 |
|
|
|
|
|
|
|
|
|
62 |
@app.post("/generate/")
|
63 |
async def generate_text(item: Item):
|
64 |
# Stream response back to the client
|
65 |
return StreamingResponse(generate_stream(item), media_type="application/x-ndjson")
|
66 |
|
|
|
|
|
67 |
# Load spaCy model
|
68 |
nlp = spacy.load("en_core_web_sm")
|
69 |
|
|
|
22 |
# Initialize the InferenceClient with your model
|
23 |
client = InferenceClient("mistralai/Mixtral-8x7B-Instruct-v0.1")
|
24 |
|
25 |
+
|
26 |
class Item(BaseModel):
|
27 |
prompt: str
|
28 |
history: list
|
29 |
system_prompt: str
|
30 |
temperature: float = 0.8
|
31 |
+
max_new_tokens: int = 12000
|
32 |
top_p: float = 0.15
|
33 |
repetition_penalty: float = 1.0
|
34 |
|
35 |
+
def format_prompt(current_prompt, history):
|
36 |
+
formatted_history = "<s>"
|
37 |
+
for entry in history:
|
38 |
+
if entry["role"] == "user":
|
39 |
+
formatted_history += f"[USER] {entry['content']} [/USER]"
|
40 |
+
elif entry["role"] == "assistant":
|
41 |
+
formatted_history += f"[ASSISTANT] {entry['content']} [/ASSISTANT]"
|
42 |
+
formatted_history += f"[USER] {current_prompt} [/USER]</s>"
|
43 |
+
return formatted_history
|
44 |
+
|
45 |
|
46 |
def generate_stream(item: Item) -> Generator[bytes, None, None]:
|
47 |
formatted_prompt = format_prompt(f"{item.system_prompt}, {item.prompt}", item.history)
|
48 |
+
# Estimate token count for the formatted_prompt
|
49 |
+
input_token_count = len(nltk.word_tokenize(formatted_prompt)) # NLTK tokenization
|
50 |
+
|
51 |
+
# Ensure total token count doesn't exceed the maximum limit
|
52 |
+
max_tokens_allowed = 32768
|
53 |
+
max_new_tokens_adjusted = max(1, min(item.max_new_tokens, max_tokens_allowed - input_token_count))
|
54 |
+
|
55 |
generate_kwargs = {
|
56 |
"temperature": item.temperature,
|
57 |
+
"max_new_tokens": max_new_tokens_adjusted,
|
58 |
"top_p": item.top_p,
|
59 |
"repetition_penalty": item.repetition_penalty,
|
60 |
"do_sample": True,
|
61 |
+
"seed": 42,
|
62 |
}
|
63 |
|
64 |
# Stream the response from the InferenceClient
|
|
|
70 |
}
|
71 |
yield json.dumps(chunk).encode("utf-8") + b"\n"
|
72 |
|
73 |
+
|
74 |
+
class SummarizeRequest(BaseModel):
|
75 |
+
text: str
|
76 |
+
|
77 |
@app.post("/generate/")
|
78 |
async def generate_text(item: Item):
|
79 |
# Stream response back to the client
|
80 |
return StreamingResponse(generate_stream(item), media_type="application/x-ndjson")
|
81 |
|
82 |
+
|
83 |
+
|
84 |
# Load spaCy model
|
85 |
nlp = spacy.load("en_core_web_sm")
|
86 |
|