Update main.py

main.py

@@ -1,85 +1,76 @@
-import os
-from typing import List, Tuple
-from fastapi import FastAPI, Form, HTTPException
-from fastapi.middleware.cors import CORSMiddleware
-from pydantic import BaseModel
-from text_generation import Client
+from fastapi import FastAPI, Request, HTTPException
+from fastapi.responses import JSONResponse, FileResponse
+from fastapi.staticfiles import StaticFiles
+from huggingface_hub import InferenceClient
+import json
 
-# Ensure the HF_TOKEN environment variable is set
-HF_TOKEN = os.environ.get("HF_TOKEN")
-if HF_TOKEN is None:
-    raise ValueError("Please set the HF_TOKEN environment variable.")
+app = FastAPI()
 
-# Model and API setup
-model_id = 'codellama/CodeLlama-34b-Instruct-hf'
-API_URL = "https://api-inference.huggingface.co/models/" + model_id
+client = InferenceClient("NousResearch/Hermes-3-Llama-3.1-8B")
 
-client = Client(
-    API_URL,
-    headers={"Authorization": f"Bearer {HF_TOKEN}"},
+SYSTEM_MESSAGE = (
+    "You are a helpful, respectful and honest assistant. Always answer as helpfully "
+    "as possible, while being safe. Your answers should not include any harmful, "
+    "unethical, racist, sexist, toxic, dangerous, or illegal content. Please ensure "
+    "that your responses are socially unbiased and positive in nature.\n\nIf a question "
+    "does not make any sense, or is not factually coherent, explain why instead of "
+    "answering something not correct. If you don't know the answer to a question, please "
+    "don't share false information."
+    "Always respond in the language of user prompt for each prompt ." 
 )
+MAX_TOKENS = 2000
+TEMPERATURE = 0.7
+TOP_P = 0.95
 
-EOS_STRING = "</s>"
-EOT_STRING = "<EOT>"
+def respond(message, history: list[tuple[str, str]]):
+    messages = [{"role": "system", "content": SYSTEM_MESSAGE}]
+    
+    for val in history:
+        if val[0]:
+            messages.append({"role": "user", "content": val[0]})
+        if val[1]:
+            messages.append({"role": "assistant", "content": val[1]})
 
-app = FastAPI()
+    messages.append({"role": "user", "content": message})
+
+    response = client.chat_completion(
+        messages,
+        max_tokens=MAX_TOKENS,
+        stream=True,
+        temperature=TEMPERATURE,
+        top_p=TOP_P,
+    )
+
+    for message in response:  # Handle regular iteration
+        yield message.choices[0].delta.content
+
+from fastapi.middleware.cors import CORSMiddleware
 
-# Allow CORS for your frontend application
 app.add_middleware(
     CORSMiddleware,
-    allow_origins=["*"],  # Change this to your frontend's URL in production
+    allow_origins=["https://artixiban-ll3.static.hf.space"],  # Allow only this origin
     allow_credentials=True,
-    allow_methods=["*"],
-    allow_headers=["*"],
+    allow_methods=["*"],  # Allow all methods (GET, POST, etc.)
+    allow_headers=["*"],  # Allow all headers
 )
 
-# Pydantic model for request body
-class ChatRequest(BaseModel):
-    prompt: str
-    history: List[Tuple[str, str]]
-
-DEFAULT_SYSTEM_PROMPT = """\
-You are a helpful, respectful and honest assistant with a deep knowledge of code and software design. Always answer as helpfully as possible, while being safe. Your answers should not include any harmful, unethical, racist, sexist, toxic, dangerous, or illegal content. Please ensure that your responses are socially unbiased and positive in nature.\n\nIf a question does not make any sense, or is not factually coherent, explain why instead of answering something not correct. If you don't know the answer to a question, please don't share false information.\
-"""
-
-def get_prompt(message: str, chat_history: List[Tuple[str, str]],
-               system_prompt: str) -> str:
-    texts = [f'<s>[INST] <<SYS>>\n{system_prompt}\n<</SYS>>\n\n']
-    do_strip = False
-    for user_input, response in chat_history:
-        user_input = user_input.strip() if do_strip else user_input
-        do_strip = True
-        texts.append(f'{user_input} [/INST] {response.strip()} </s><s>[INST] ')
-    message = message.strip() if do_strip else message
-    texts.append(f'{message} [/INST]')
-    return ''.join(texts)
-
 @app.post("/generate/")
-async def generate_response(prompt: str = Form(...), history: str = Form(...)):
-    try:
-        chat_history = eval(history)  # Convert history string back to list
-        system_prompt = DEFAULT_SYSTEM_PROMPT
-        message = prompt
+async def generate(request: Request):
+    allowed_origin = "https://artixiban-ll3.static.hf.space"
+    origin = request.headers.get("origin")
+    if origin != allowed_origin:
+        raise HTTPException(status_code=403, detail="Origin not allowed")
+    form = await request.form()
+    prompt = form.get("prompt")
+    history = json.loads(form.get("history", "[]"))  # Default to empty history
 
-        prompt_text = get_prompt(message, chat_history, system_prompt)
+    if not prompt:
+        raise HTTPException(status_code=400, detail="Prompt is required")
 
-        generate_kwargs = dict(
-            max_new_tokens=1024,
-            do_sample=True,
-            top_p=0.9,
-            top_k=50,
-            temperature=0.1,
-        )
-        
-        stream = client.generate_stream(prompt_text, **generate_kwargs)
-        output = ""
-        for response in stream:
-            if any([end_token in response.token.text for end_token in [EOS_STRING, EOT_STRING]]):
-                break
-            else:
-                output += response.token.text
+    response_generator = respond(prompt, history)
+    final_response = ""
+    for part in response_generator:
+        final_response += part
 
-        return {"response": output}
+    return JSONResponse(content={"response": final_response})
 
-    except Exception as e:
-        raise HTTPException(status_code=500, detail=str(e))