File size: 2,860 Bytes
a0957ab
3ff3d0b
 
 
 
 
 
 
a0957ab
3ff3d0b
a0957ab
3ff3d0b
 
3ccd4b4
3ff3d0b
 
 
 
 
 
 
 
3ccd4b4
3ff3d0b
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
8d74a81
3ff3d0b
 
3ccd4b4
3ff3d0b
 
 
 
 
8d74a81
3ff3d0b
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
3ccd4b4
 
3ff3d0b
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
#from peft import AutoPeftModelForCausalLM
import asyncio
from fastapi import FastAPI
import uvicorn
import os
import nest_asyncio
import socket
import json
#from transformers import pipeline, AutoTokenizer
import time
#import torch
import requests
app = FastAPI()
"""

model_ask = AutoPeftModelForCausalLM.from_pretrained(

    "robertthecreator/assistant-chatbot",

    device_map="auto",

    torch_dtype=torch.float16

)

tokenizer_ask = AutoTokenizer.from_pretrained("robertthecreator/assistant-chatbot")

tokenizer_ask.pad_token = tokenizer_ask.eos_token

model_ask.eval()

"""
api_key = os.getenv("TAVILY_API_KEY", "default-key")

def get_from_web(query, k=3):
    context = ""
    response = requests.post(
        "https://api.tavily.com/search",
        headers={"Content-Type": "application/json"},
        json={"query": query, "api_key": api_key}
    )
    data = response.json()
    for result in data["results"]:
        print(result["content"])
        context += result["content"] + "\n"
    return context


@app.post("/ask")
def ask(text: str):
    global context
    print("starting...")
    context = """You're an AI assistant designed to help me with every day task, productivity and engineering. Youre name is JARVIS. Use the context you have to answer any questions."""
    #context = get_from_web(text) + "\n" + context
    print(context)
    formatted_prompt = f"""CONTEXT: {context}\nPROMPTER: {text}\nASSISTANT:"""
    """

    inputs = tokenizer_ask(formatted_prompt, return_tensors="pt").to(model_ask.device)



    with torch.no_grad():

        outputs = model_ask.generate(

            **inputs,

            max_new_tokens=100,

            temperature=0.7,

            top_p=0.9,

            top_k=50,

            repetition_penalty=1.2,

            do_sample=True,

            pad_token_id=tokenizer_ask.eos_token_id,

            eos_token_id=tokenizer_ask.eos_token_id,

            early_stopping=True

        )





    response = tokenizer_ask.decode(outputs[0], skip_special_tokens=True)

    print(response)



    if "ASSISTANT:" in response:

        response = response.split("ASSISTANT:")[-1].strip()



    response = response.split('\r')[0]

    response = response.split('PROMPTER')[0]

    response = response.split('ROLE:')[0]

    response = response.split('<|user|>')[0]

    response = response.split('User:')[0]

    """
    response = "This is a placeholder response."
    return response

async def start():
    s = socket.socket()
    s.bind(('', 0))
    port = s.getsockname()[1]
    s.close()

    config = uvicorn.Config(
        app,
        host="0.0.0.0",
        port=port,
        log_level="info"
    )
    server = uvicorn.Server(config)

    await server.serve()

uvicorn.run(app, host="0.0.0.0", port=8000, log_level="info")