Spaces:
Runtime error
Runtime error
File size: 2,780 Bytes
531d852 | 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 | 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=250,
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]
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") |