Update main.py

main.py

@@ -1,51 +1,56 @@
 from fastapi import FastAPI
 from fastapi.staticfiles import StaticFiles
 from fastapi.responses import FileResponse
-from transformers import T5Tokenizer, T5ForCausalLM, Trainer, TrainingArguments
+from transformers import BertTokenizer, BertForMaskedLM, Trainer, TrainingArguments
 
 app = FastAPI()
 
 # Initialize the tokenizer and model
-tokenizer = T5Tokenizer.from_pretrained("t5-base")
-model = T5ForCausalLM.from_pretrained("t5-base")
+tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")
+model = BertForMaskedLM.from_pretrained("bert-base-uncased")
 
+# Prepare the training data
 with open("cyberpunk_lore.txt", "r") as f:
-    dataset = f.read()
+    train_data = f.read()
+train_data = train_data.split("\n")
+train_data = [tokenizer.encode(text, return_tensors="pt") for text in train_data]
 
-# Tokenize the dataset
-input_ids = tokenizer.batch_encode_plus(dataset, return_tensors="pt")["input_ids"]
-
-# Set up training arguments
+# Define the training arguments
 training_args = TrainingArguments(
-    output_dir='./results',
-    overwrite_output_dir=True,
-    num_train_epochs=5,
-    per_device_train_batch_size=1,
+    output_dir="./results",
+    per_device_train_batch_size=16,
     save_steps=10_000,
     save_total_limit=2,
 )
 
-# Create a Trainer
+# Create the trainer
 trainer = Trainer(
     model=model,
     args=training_args,
-    train_dataset=input_ids,
-    eval_dataset=input_ids
+    train_dataset=train_data,
+    eval_dataset=train_data,
 )
 
-# Fine-tune the model
+# Start the training
 trainer.train()
 
-# Create the inference pipeline
-pipe_flan = pipeline("text2text-generation", model=model)
+# Save the fine-tuned model
+trainer.save_model('./results')
 
-@app.get("/infer_t5")
-def t5(input):
-    output = pipe_flan(input)
-    return {"output": output[0]["generated_text"]}
+# Load the fine-tuned model
+model = trainer.get_model()
 
-app.mount("/", StaticFiles(directory="static", html=True), name="static")
+# Create the inference endpoint
+@app.post("/infer")
+def infer(input: str):
+    input_ids = tokenizer.encode(input, return_tensors="pt")
+    output = model(input_ids)[0]
+    return {"output": output}
 
+@app.get("/")
+def index() -> FileResponse:
+    return FileResponse(path="/app/static/index.html", media_type="text/html")
+    
 @app.get("/")
 def index() -> FileResponse:
     return FileResponse(path="/app/static/index.html", media_type="text/html")